Learning Interpretable Features in Audio Latent Spaces via Sparse Autoencoders

AI summary

Overview

Research area: Mechanistic interpretability of generative models, applied to audio and music generation; sparse autoencoders (SAEs), latent-space probing, and controllable audio synthesis.

Technical level: Intermediate. The paper assumes familiarity with sparse autoencoders, linear probes, and audio autoencoder latents, but its framing is accessible to readers with a general machine-learning background.

Scope: The paper trains sparse autoencoders on the latent spaces of continuous and discrete audio autoencoders, learns linear mappings from those sparse features to discretized acoustic properties (pitch, amplitude, timbre), and uses those mappings both to edit audio and to trace how acoustic properties emerge during generation.

What This Paper Is About

Sparse autoencoders have worked well for finding interpretable features inside language models, but audio is a harder case: it is dense rather than token-like, it must be compressed through an autoencoder before modeling, and no audio model can yet automatically describe what an SAE feature means the way a language model can summarize its own perturbations. The paper's goal is to close that interpretation gap by linking SAE features from audio generation latents to acoustic concepts a human can name and measure, and then using those links to control and analyze music generation.

Key Contributions

1. A framework for interpreting audio generative models through acoustic concepts. SAEs are trained on audio autoencoder latents, and linear probes map the resulting sparse features to discretized pitch, amplitude, and timbre. The authors argue the success of these linear mappings is itself evidence that SAE features already encode acoustic properties in a near-linear fashion.

2. An SAE architecture modified for audio latents. A RMS normalization layer is inserted after the ReLU activation to keep activation magnitudes consistent and, as the authors report empirically, to prevent out-of-distribution artifacts when manipulating features.

3. Validation across continuous and discrete latent spaces. The method is tested on the continuous VAE space of Stable Audio Open and DiffRhythm and on the discrete latent spaces of EnCodec and WavTokenizer, showing the approach is not tied to one tokenization scheme.

4. Use of the learned mappings to trace the generation process. By decomposing intermediate latents of DiffRhythm, a state-of-the-art text-to-music model, the authors track when pitch, timbre, and loudness emerge across inference steps, and they demonstrate targeted audio edits via "control vectors."

Main Findings

• Sparsity varies systematically by latent space. DiffRhythm's SAEs reach sparsity ratios of 0.65 to 0.98; WavTokenizer produces the sparsest representations at 0.993 to 0.999, which the authors suggest means its discrete tokens already encode highly disentangled features; EnCodec shows the widest range at 0.55 to 0.95. Across all models, larger hidden dimensions consistently improve reconstruction quality.

• Acoustic properties form a hierarchy of linear decodability. Pitch is the most linearly separable property, with probe accuracy of 0.75 to 0.87, and stays stable across all sparsity levels, suggesting it reflects fundamental frequency encoding. Timbre remains challenging across all models at 0.17 to 0.46. EnCodec reaches 0.56 to 0.63 on loudness, while DiffRhythm and WavTokenizer reach only 0.17 to 0.49.

• Targeted control vectors produce isolated edits. On a chordal CocoChorales sample encoded with EnCodec and the highest-sparsity SAE (hidden_dim=16384, lambda=0.1), control vectors targeting pitch (MIDI C5), timbre (spectral centroid class 17), and loudness (class 2) at strengths alpha in {1, 10, 20, 30} changed the targeted attribute while non-targeted properties were largely preserved.

• DiffRhythm establishes fundamental frequency before texture and dynamics. Analyzing a 95-second generated segment spanning a verse and a chorus over 32 inference steps, with 500 prompts sampled from MusicCaps, pitch converges first (around step 21), timbre follows, and loudness converges last and remains unresolved by the final step. The authors describe this as a coarse-to-fine progression.

• Probe setup details. The probes predict pitch with 66 bins spanning the pitch range present in the dataset, loudness with 20 bins, and timbre with 20 bins. Pitch uses logarithmic bins aligned with MIDI note numbers, while amplitude and timbre use linear bins.

Methodology in Plain English

The researchers take pretrained audio autoencoders, whose job is to compress audio into a latent representation, and train sparse autoencoders on top of those latent representations. An SAE is a small network that tries to rebuild its input using only a few active hidden units, which pushes it toward finding clean, separable directions in the input space. The authors deviate from the standard recipe by adding RMS normalization after the ReLU activation so that feature magnitudes stay consistent — a change they found prevents strange artifacts when features are later edited.

Once the SAEs are trained, the authors extract measurable acoustic quantities from the audio itself: pitch using CREPE, amplitude using windowed RMS energy from librosa, and timbre using windowed spectral centroid from librosa as a simplified proxy. Each of these continuous curves is then chopped into discrete classes, and a simple linear classifier is trained to predict the class from the SAE features. Because the classifier is linear, the weight from a given SAE feature to a given acoustic class directly shows how much that feature contributes — and reversing that logic lets the authors add a scaled weight vector to the SAE features to push audio toward a chosen acoustic class. The edited features are then decoded back through the SAE and the audio decoder to produce modified audio.

To study generation, the authors run DiffRhythm for 32 inference steps and, at each step, decompose the intermediate latent through the same SAE and probes, average across frames, and measure how far each acoustic property's distribution has traveled from the first step toward the final one.

Why This Matters

Impact on research. The paper extends mechanistic interpretability, which has largely been developed on language models, into a modality where the standard toolkit does not transfer cleanly. It shows that SAE features in audio latents align with measurable acoustic concepts, and it supplies an audit trail for how those concepts arise inside a diffusion-style music generator — something prior work could not characterize automatically. The authors note their framework is only demonstrated on audio but believe it generalizes to other generative models operating in learned latent spaces, including images and video.

Real-world applications:

• Music production tools. Control vectors that shift pitch, timbre, or loudness in isolation point toward DAW-style tools where a producer edits a generated stem along interpretable axes rather than regenerating from scratch.

• Debugging and quality control for generative audio systems. Generation-progress analysis can reveal whether a model has locked in a property (like pitch) before it has resolved another (like loudness), which is useful for diagnosing failures in long-form synthesis.

• Transparency and model auditing. Mapping latent directions to named acoustic concepts gives a way to describe what a black-box music generator is doing internally, which matters for content provenance and safety review.

• Improved conditioning interfaces. Because the mappings are bidirectional, acoustic concepts can be decomposed into specific SAE features, potentially enabling more precise text or attribute conditioning in music generation products.

Industry relevance. The work was done in part during an internship at Smule Labs, with co-authors from Smule, UC San Diego, and Stanford University, which places it directly adjacent to consumer music and singing applications. Controllable, interpretable editing of generated audio is a practical requirement for any product that lets users shape music rather than just prompt it.

Future Directions

• Applying the framework to other generative architectures. The authors name RAVE, ACE-Step, and AudioLDM as targets for extending the method beyond the models tested here.

• Probing richer audio features. Beyond pitch, amplitude, and timbre, they plan probes for rhythm, harmony, and instrument identity.

• Guiding generation directly at inference time. The stated goal is to use interpretable features to steer generation behavior mid-process, aiming for fine-grained control over specific attributes while maintaining generation quality.

• Extending beyond audio. The authors suggest the same latent-space interpretation approach could be applied to visual and video generative models, though they emphasize this has not yet been tested.

Target Audience

This paper is most useful to mechanistic interpretability researchers interested in moving beyond language models, audio and music generation researchers who want controllable and inspectable models, and machine-learning engineers building production music tools who need interpretable knobs rather than opaque prompts. Readers with a background in representation learning or audio signal processing will get the most from it; those unfamiliar with sparse autoencoders will need to pick up that background first.

Authors’ abstract

While sparse autoencoders (SAEs) successfully extract interpretable features from language models, applying them to audio generation faces unique challenges: audio's dense nature requires compression that obscures semantic meaning, and automatic feature characterization remains limited. We propose a framework for interpreting audio generative models by mapping their latent representations to human-interpretable acoustic concepts. We train SAEs on audio autoencoder latents, then learn linear mappings from SAE features to discretized acoustic properties (pitch, amplitude, and timbre). This enables both controllable manipulation and analysis of the AI music generation process, revealing how acoustic properties emerge during synthesis. We validate our approach on continuous (DiffRhythm-VAE) and discrete (EnCodec, WavTokenizer) audio latent spaces, and analyze DiffRhythm, a state-of-the-art text-to-music model, to demonstrate how pitch, timbre, and loudness evolve throughout generation. While our work is only done on audio modality, our framework can be extended to interpretable analysis of visual latent space generation models.