Note
Code for the ICLR 2026 paper Discovering and Steering Interpretable Concepts in Large Generative Music Models. You can find examples in our interactive dashboard here.
We introduce a method to discover musical concepts by training sparse autoencoders (SAEs) on the residual stream activations of a transformer-based music generation model (MusicGen), and then to automatically label such concepts.
Create the environment as shown below (we recommend using uv):
conda create -n musicdiscovery python=3.10
conda activate musicdiscovery
conda install -c conda-forge ffmpeg
pip install uv
uv pip sync requirements.txtIf you plan to run Essentia labeling, download required TensorFlow models once:
./scripts/dl_essentia.shThis codebase will currently support the following datasets:
- MusicSet: 10-second audio clips (MusicSet corpus) configured in
conf/dataset/musicset.yaml. This is the dataset we use in the paper's results, but it has since unfortunately been taken down. - FMA: Clips from the Free Music Archive (small/large) configured in
conf/dataset/fma-small.yamlandconf/dataset/fma-large.yaml
Override the default cache location for these with general.cache_dir in the configuration if needed.
Important
We use Hydra to configure everything. The configuration can be found in conf/config.yaml, with specific sub-configs in sub-directories of conf/.
The default config has model=musicgen-small and dataset=fma-small.
This step is optional but recommended for faster subsequent training runs. It caches the activations of a MusicGen model's layer for the specified dataset.
python3 cache_activations_with_musicgen.py \
model=musicgen-small \
dataset=fma-small \
saeconfig=musicgen-main \
cacheconfig=musicgen-main \
saeconfig.hook_layer=21This step trains a sparse autoencoder, either by obtaining activations as needed, or from the cached activations from the previous step.
python3 train_sae_on_musicgen.py \
model=musicgen-small \
dataset=fma-small \
saeconfig=musicgen-main \
saeconfig.use_cached_activations=true \
saeconfig.cached_activations_path="$(pwd)/outputs/<path-to-cached-activations>" \
saeconfig.hook_layer=21This will apply the trained SAE to the validation set of the given dataset, and write the encoded features to disk. The config's evaluation.checkpoints_dir should point to the directory where the SAE model checkpoints are saved.
python3 extract_concepts.py \
dataset=fma-small \
evaluation.checkpoints_dir="$(pwd)/outputs/<path-to-checkpoints-dir>"This writes validation-features.pt (by default) under each discovered checkpoint subdirectory.
Labeling the concepts requires a few extra steps.
- collect max-activating example candidates:
python3 scripts/collect_samples.py \
--features-dir "$(pwd)/outputs/<path-to-checkpoints-dir>" \
--features-file validation-features.ptNote the --dataset argument (you may need to set this).
- label the concepts using Essentia (you may need to first run
./scripts/dl_essentia.sh):
python3 scripts/label_concepts-essentia.py \
--input_dir concepts_validation-features \
--set validation-features_pct-selective \
--output_file data/labels_essentia.jsonl \
--n_workers 16- label the concepts using Gemini.
Set up credentials for google.genai (for example via GEMINI_API_KEY or equivalent Google AI auth config), then:
python3 scripts/label_concepts-gemini.py \
--features_path concepts_validation-features \
--output_file data/concepts_gemini.jsonl \
--model gemini-3-flash-preview \
--jobs 32Compute CLAP-based embeddings and coherence:
python3 scripts/concept_embeddings.py \
--input_dir concepts_validation-features \
--num_workers 16
python3 scripts/concept_coherence.py \
--input_dir concepts_validation-features \
--num_workers 16
# Download the CLAP checkpoint
wget https://huggingface.co/lukewys/laion_clap/resolve/main/music_audioset_epoch_15_esc_90.14.pt -O data/music_audioset_epoch_15_esc_90.14.pt
python3 scripts/compute_clapscore.py \
--input_dir concepts_validation-features \
--essentia_file data/labels_essentia.jsonl \
--gemini_file data/concepts_gemini.jsonl \
--outfile data/scores_clap.jsonIf you switch input directories between runs, remove the old cached file lists so scripts rescan from the right source:
rm -f feature_dirs.json feature_dirs_cache.jsonIf you want to quickly compress the audio files for deployment (e.g. in an interactive dashboard, like ours):
python3 scripts/compress_tree_of_audio.py \
--input_dir concepts_validation-features \
--output_dir dashboard_audio \
--num_workers 16Additional helper scripts in scripts/ include checkpoint gathering and more:
scripts/gather_checkpoints.py
First, run a bootstrap pass to steer all available features and identify the most steerable ones. Omit --features to steer all features. Note that PLACEHOLDER is a literal token in the checkpoint path — the script replaces it with each SAE name automatically:
python3 steering_concepts.py \
--checkpoint /path/to/PLACEHOLDER/facebook/musicgen-large/cfg.json \
--dataset concepts_validation-features \
--output results/bootstrapEvaluate the bootstrap audio to identify the most steerable features:
python3 evaluate_steering.py \
--results results/bootstrap/musicgen-large \
--dataset concepts_validation-features \
--output results/bootstrap/musicgen-large/evaluation_results.jsonRank features by steering improvement and export top_steerable_features.json:
python3 scripts/analyze_steering_results.py \
--results results/bootstrap/musicgen-large/evaluation_results.json \
--strength 1.0 \
--baseline 3This writes top_steerable_features.json and steering_improvement_distributions.png alongside the results file.
Run a targeted pass on the top features using --features, optionally sweeping multiple strengths:
python3 steering_concepts.py \
--checkpoint /path/to/PLACEHOLDER/facebook/musicgen-large/cfg.json \
--dataset concepts_validation-features \
--output results/steered \
--features results/bootstrap/musicgen-large/top_steerable_features.json \
--strengths 0.5 1.0 2.0Evaluate the generated audio against feature examples using CLAP embeddings:
python3 evaluate_steering.py \
--results results/steered/musicgen-large \
--dataset concepts_validation-features \
--output results/steered/musicgen-large/evaluation_results.jsonCompute steering improvement metrics, rank features, and export a LaTeX table:
python3 scripts/analyze_steering_results.py \
--results results/steered/musicgen-large/evaluation_results.json \
--strength 1.0 \
--baseline 3This writes top_steerable_features.json and steering_improvement_distributions.png alongside the results file.
Collect the top-activating feature example audios for the top steerable features:
python3 scripts/collect_concept_samples.py \
--json results/steered/musicgen-large/top_steerable_features.json \
--data-dir concepts_validation-features \
--output-dir collected_featuresCollect the corresponding baseline and steered audio files into a flat directory:
python3 scripts/collect_steered_samples.py \
--features results/steered/musicgen-large/top_steerable_features.json \
--results results/steered/musicgen-large \
--output collected_steered \
--strength 1.0Jupyter notebooks under analysis/ run the visualizations and other analyses described in the paper. They include:
analysis/plot_feature_layer_analysis.ipynb: Analyzes feature counts and layer activations.analysis/plot_clap_scores_analysis.ipynb: Analyzes CLAP scores across model variants and layers.analysis/utils.py: Supporting functions for the analysis notebooks.
We currently host SAE checkpoints on AWS S3 (cfg = config, w = weights, sp = sparsity).
| exp | k | L2 | L6 | L12 | L18 | L22 |
|---|---|---|---|---|---|---|
| 4 | 32 | cfg w sp | cfg w sp | cfg w sp | cfg w sp | cfg w sp |
| 4 | 100 | cfg w sp | cfg w sp | cfg w sp | cfg w sp | cfg w sp |
| 32 | 32 | cfg w sp | cfg w sp | cfg w sp | cfg w sp | cfg w sp |
| 32 | 100 | cfg w sp | cfg w sp | cfg w sp | cfg w sp | cfg w sp |
| exp | k | L2 | L12 | L24 | L36 | L46 |
|---|---|---|---|---|---|---|
| 4 | 32 | cfg w sp | cfg w sp | cfg w sp | cfg w sp | cfg w sp |
| 4 | 100 | cfg w sp | cfg w sp | cfg w sp | cfg w sp | cfg w sp |
| 32 | 32 | cfg w sp | cfg w sp | cfg w sp | cfg w sp | cfg w sp |
| 32 | 100 | cfg w sp | cfg w sp | cfg w sp | cfg w sp | cfg w sp |
To download them all:
BASE="https://music-discovery-sae-checkpoints.s3.amazonaws.com"
download_bundle() {
for exp in 4 32; do
for k in 32 100; do
for layer in "$@"; do
prefix="sae-${exp}_k_${k}_${layer}/facebook/${MODEL}"
outdir="sae-${exp}_k_${k}_layer_$((layer + 1))/${MODEL}"
mkdir -p "$outdir"
echo "Downloading ${prefix}"
curl -fL -o "${outdir}/cfg.json" "${BASE}/${prefix}/cfg.json"
curl -fL -o "${outdir}/sae_weights.safetensors" "${BASE}/${prefix}/sae_weights.safetensors"
curl -fL -o "${outdir}/sparsity.safetensors" "${BASE}/${prefix}/sparsity.safetensors"
done
done
done
}
MODEL="musicgen-small"; download_bundle 1 5 11 17 21
MODEL="musicgen-large"; download_bundle 1 11 23 35 45If this work was useful, please cite the following paper:
@article{singh2026discovering,
title={Discovering and Steering Interpretable Concepts in Large Generative Music Models},
author={Singh, Nikhil and Cherep, Manuel and Maes, Pattie},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://arxiv.org/abs/2505.18186},
}We gratefully acknowledge AWS for providing computational resources through the AWS Research and Engineering Studio (RES) pilot program. We thank Brian McCarthy and Jared Novotny for their valuable support, and Peter Fisher (MIT ORCD) for making this collaboration happen. The authors acknowledge the MIT Office of Research Computing and Data for providing high performance computing resources that have contributed to the research results reported within this paper. MC is supported by a fellowship from “la Caixa” Foundation (ID 100010434) with code LCF/BQ/EU23/12010079. We also thank Michael Casey, Anna Huang, and the HAI-Res lab for supportive comments.