Technical Summary: Residualized Temporal Sparse Autoencoders for Interpreting Diffusion Models

Problem Statement

Text-to-image diffusion models generate images through an iterative denoising process, causing internal neural layers to produce trajectories of activations rather than single static representations. While Sparse Autoencoders (SAEs) have recently been applied to decompose diffusion activations into interpretable feature directions, existing approaches typically analyze activations at isolated timesteps, condition on time, or apply SAEs to architectures like Diffusion Transformers.

A critical limitation of these existing methods is that they do not explicitly account for the strong temporal correlation inherent in diffusion activations. Activations at adjacent timesteps are highly correlated, with much of the signal at one step being linearly predictable from the previous step. Training an SAE directly on raw, temporally concatenated activations risks wasting model capacity on learning these redundant, linearly predictable dynamics rather than capturing the underlying sparse structure of the features.

Methodology

The authors propose Residualized Temporal SAEs, a framework designed to learn sparse features from activation trajectories by explicitly separating linearly predictable dynamics from residual structure.

1. Data Collection and Preprocessing

• Source: Activations are collected from the conditional branch of the U-Net in Stable Diffusion 1.5 during classifier-free guidance. Specifically, the residual contribution of the mid-block cross-attention layer is extracted.
• Trajectory Construction: For a generated image, activations are subsampled every $\ell$ denoising steps (e.g., $\ell=10$) to form a sequence of activation tensors $\{a(\tau_i)\}$.

2. Linear Residualization

To remove temporal redundancy, the method fits linear predictors between consecutive normalized activation blocks:

• Normalization: Activations are normalized per timestep using training-set statistics.
• Ridge Regression: For each adjacent transition $i-1 \to i$, a ridge regressor $(W_i, b_i)$ is fitted to predict the current activation $\bar{a}_{q,i}$ from the previous one $\bar{a}_{q,i-1}$.
• Residual Calculation: The residual $r_{q,i} = \bar{a}_{q,i} - (W_i \bar{a}_{q,i-1} + b_i)$ captures the component of the activation not linearly predictable from the previous step.

3. SAE Training on Residualized Trajectories

The input to the SAE is a concatenated vector $z_q$ consisting of the initial activation and the subsequent residuals:
$$z_q = [\bar{a}_{q,0}, r_{q,1}, \dots, r_{q,T-1}]$$
A BatchTopK SAE is trained on this residualized representation. The objective is to reconstruct the normalized input while enforcing sparsity (retaining only the top $K$ active latents per batch). This encourages the sparse latents to capture structure that exists beyond the linearly predictable temporal drift.

4. Activation-Space Interpretation

A key component of the method is mapping the learned features back to the original activation space to analyze them as trajectories:

1. Decoder Directions: The SAE decoder vectors are initially in the normalized residualized space.
2. Un-normalization: Scales are applied to return to the normalized residual space.
3. Linear Dynamics Integration: The residual directions are mapped back through the fitted linear predictors ($W_i$) to reconstruct the full trajectory direction in normalized activation space.
4. Final Mapping: The result is un-normalized to obtain the decoder trajectory $\{\phi_{k,i}\}$ in the original activation units. This allows each latent to be analyzed as a feature trajectory evolving over denoising time.

Key Contributions

1. Residualized Temporal SAEs: The introduction of a method that separates diffusion activation trajectories into linearly predictable components and residuals, training SAEs specifically on the latter to avoid modeling redundant temporal dynamics.
2. Activation-Space Decoder Interpretation: A derivation that maps residualized decoder directions back into the original activation space, enabling the analysis of each latent as a spatiotemporal feature trajectory.
3. Empirical Characterization: A comprehensive evaluation on Stable Diffusion 1.5 demonstrating that these representations capture temporally structured, spatially localized directions with semantically meaningful generative effects.

Results and Evaluation

Reconstruction Performance

• Superiority of Residualization: Under a fixed trajectory-level sparsity budget, residualized variants consistently outperform non-residualized counterparts (both concatenated and timestep-wise) in Mean Squared Error (MSE) and Explained Variance (EV).
• Efficiency: The Residualized Concatenated model achieves competitive reconstruction performance while using a single sparse code for the entire trajectory, whereas timestep-wise baselines require distributed codes.
• Baseline Comparison: The method outperforms a Matryoshka SAE baseline (which uses non-residualized, component-normalized trajectories) in reconstruction quality.

Spatiotemporal Structure Analysis

• Temporal Specialization: Latents exhibit distinct temporal profiles. "Early" features tend to have broad temporal support and high cross-timestep self-similarity, suggesting they capture persistent structural concepts. "Late" features are more timestep-specific and less temporally stable.
• Spatial Localization:
  • Early features show broader positive spatial support (higher spatial entropy), often covering diffuse or border-like regions.
  • Middle and Late features are more spatially concentrated (lower entropy), localizing to specific object-like areas or salient regions.

Steering Experiments

• Single-Feature Steering: Intervening with specific latents produces visible generative effects. Early features affect global structure, while middle features can localize changes (e.g., adding jewelry to a torso). Late features showed less clear semantic effects in some examples.
• Feature Transfer: Using the RIEBench dataset, the authors performed source-to-target feature transfer. The Residualized Concatenated model achieved the best balance between editing efficiency (improving CLIP similarity to the target prompt) and preserving the source image structure (minimizing LPIPS distance). It outperformed non-residualized and timestep-wise variants in edit efficiency.

Significance and Claims

The paper claims that Residualized Temporal SAEs provide a useful framework for studying temporally structured diffusion activations. By explicitly removing linearly predictable drift, the method allows sparse latents to focus on the non-redundant, semantically meaningful structure of the generation process.

The authors emphasize that their approach differs from prior work by learning sparse dictionaries over residualized activation trajectories rather than isolated timesteps or non-residualized temporal representations. The ability to map these latents back into activation space as trajectories enables a deeper understanding of when and where specific features are expressed during the denoising process.

Limitations noted by the authors:

• The experiments focus on a single U-Net activation source (mid-block cross-attention) in Stable Diffusion 1.5.
• Steering experiments are qualitative; the paper does not claim optimized image-editing performance but rather uses steering to validate the interpretability and causal influence of the learned features.