ELROND: Exploring and decomposing intrinsic capabilities of diffusion models

The gist A single text prompt passed to a diffusion model yields a wide range of visual outputs determined solely by a stochastic process, leaving users with no direct control over which semantic variations appear exploring this range is difficult because random search offers no guarantee of covering it and prompt editing is coarse.

How it works

The paper introduces ELROND, a method for recovering and decomposing the tangent space of the generative manifold that a diffusion model has learned. This recovery involves collecting gradients obtained by backpropagating the differences between stochastic realizations of a fixed prompt and decomposing them into interpretable directions using Principal Component Analysis or a Sparse Autoencoder. The core idea is to view the model through its family of conditional distributions, referred to as the model’s generative manifold, where each point is a distribution of images.

The process for extracting semantic directions involves several steps:

1. Sample a dataset of N images, X = x(1),..., x(N) from the prompt.

2. For a pair of distinct images (x(i), x(j)), construct a reconstruction objective that encourages the model’s expected output to move from the characteristics of one realization toward those of the other.

3. Backpropagate this objective to the conditioning embedding, obtaining a direction in the conditioning space that corresponds to a local displacement on the generative manifold.

4. Repeating this procedure over many pairs provides a collection of local directions that span the tangent space at the distribution induced by the prompt.

Discovering Latent Directions

The raw gradients in G often contain entangled information about several semantic changes together with unimportant noise, necessitating decomposition into disentangled, semantic directions. Two primary approaches are employed for this decomposition:

1. Principal Component Analysis (PCA): PCA is performed on G, and the resulting latent directions are defined as the principal components (eigenvectors) vk. Noise is filtered by retaining only components with explained variance above a threshold τ.

2. Sparse Autoencoder Decomposition (SAE): A Top-k Sparse Autoencoder is trained on G, where an encoder projects the input vector v to a higher-dimensional latent space h using a linear transformation followed by a TopK activation function. The decoder then reconstructs the input from this sparse latent code h, and the decoder weights corresponding to these features represent the disentangled semantic directions.

Semantic Steering

Once a discovered latent direction d is obtained, it can be used to precisely steer the generation process by modifying the concept token embedding. The modified embedding is calculated as e' = e + λ · d, where λ is a strength parameter. This method scales naturally to complex prompts with multiple subjects, as directions are derived relative to specific tokens.

Validation and Results

The framework was validated first in a controlled setting where the tangent space is known, showing that the recovered gradients recover the ground-truth subspace. On large-scale text-to-image models, it was shown that the model’s sensitivity to its conditioning is concentrated in the recovered subspaces. The method enables broader exploration of a model’s capabilities than methods relying on external representations and mitigates mode collapse in distilled models without retraining.

The comparison between PCA and SAE decomposition showed that PCA consistently yields the highest Recall scores, while SAE achieves higher Precision because it adapts to the non-uniform density of the semantic gradients. Furthermore, ELROND successfully mitigates mode collapse in distilled models by using an SAE variant that captures the probability density of the concept manifold, leading to a lower FID than both SliderSpace and the base distilled model.

The user study confirmed that 40 human evaluators significantly preferred ELROND in terms of diversity and creativity, while on image fidelity, our method’s win rate falls within the confidence interval of being statistically equivalent to the unperturbed model. The method also demonstrates token-level composability, enabling independent steering of distinct concepts through token-level control.

The paper concludes that ELROND offers a geometry-grounded account of a model’s generative variety, derived from the model itself rather than imposed upon it. The effectiveness is demonstrated across various architectures, including SDXL and Flux Schnell. The method's ability to capture the intrinsic semantic manifold of a given concept is robust to initial noise sampling.

Limitations

ELROND has limitations, including the computational expense of data collection due to backpropagating through the diffusion model. Additionally, the concept-level analysis relies on the assumption that token embeddings are localized, meaning perturbing one token can affect other tokens. Some directions discovered by either SAE or PCA may not admit a clear interpretation. The contextual dependency of text embeddings can result in partial semantic leaking when applying latent directions to specific subjects.

Loss Functions for Semantic Gradient Calculation

To assess the efficacy of different loss functions, four formulations were compared for calculating the semantic gradient gi,j between a source image x(i) and a target image x(j). The standard Mean Squared Error (MSE) loss consistently outperformed other formulations across all three metrics, achieving the lowest DreamSim distance (0.4114), the highest CLIP similarity (0.8037), and the highest DINO similarity (0.5073).

The paper demonstrates that standard MSE loss is crucial for extracting accurate semantic directions in the diffusion embedding space. The choice of loss function significantly influences the quality and steerability of the resulting semantic directions, with standard MSE being the optimal approach.

Comparison with SliderSpace

When benchmarked under an equal computational budget, ELROND maintains competitive diversity and text alignment even when gradient sampling is drastically reduced. While SliderSpace achieves high Precision scores for certain concepts, ELROND lags behind in Recall, confirming that deriving directions from the model’s own gradients is better suited for uncovering the full latent diversity of the model than when constrained by external proxies.

The paper shows that simple linear interpolation in the text conditioning space yields highly non-linear paths in the visual embedding space, necessitating a gradient-based approach like ELROND. The research provides a geometry-grounded account of a model’s generative variety, derived from the model itself rather than imposed upon it.

The paper is 2602.10216.

1. Make diffusion models more controllable by recovering their generative manifold structure, as ELROND recovers its tangent space in the conditioning space. This allows for systematic exploration of model capabilities beyond coarse prompt editing or random search.

2. Enable precise, concept-level steering within generated images by injecting recovered latent directions into token embeddings, as this method allows for token-level composability. This enables users to apply distinct directions to different subjects simultaneously with limited cross-concept interference.

3. Mitigate mode collapse in distilled models without retraining by using learned structures to restore diversity, as the SAE variant of ELROND can achieve a lower FID than both the base distilled model and SliderSpace, because the SAE dictionary adapts to the non-uniform density of the semantic gradients.

4. Improve attribute control fidelity by utilizing optimized loss functions for gradient calculation, as standard MSE loss consistently outperforms other formulations across all three metrics for achieving the highest DreamSim distance (0.4114) and maintaining high CLIP similarity.

5. Enhance model generalization across different architectures by extending ELROND to transformer-based models like Flux Schnell, as the method successfully identifies semantically rich and steerable directions within the conditioning space of the transformer architecture.

A single text prompt passed to a diffusion model yields a wide range of visual outputs determined solely by a stochastic process, leaving users with no direct control over which semantic variations appear. Exploring this range is difficult: random search offers no guarantee of covering it, while prompt editing is coarse, as even a small change in wording can substantially alter the generated image. We argue that systematic exploration instead requires recovering how the model itself organizes the conditional distributions it can produce. We formalize this structure as a generative manifold, and present ELROND, a method for recovering its tangent space at a given conditioning. To that end, we collect gradients obtained by backpropagating the differences between stochastic realizations of a fixed prompt, and decompose them into interpretable directions using Principal Component Analysis or a Sparse Autoencoder. We show that our method recovers this subspace accurately in a controlled setting and validate it behaviorally on large-scale models. We also demonstrate that recovered structure enables broader exploration of the model's capabilities than methods relying on external representations, and mitigates mode collapse in distilled models without retraining.