Test-Time Conditioning with Representation-Aligned Visual Features

TL;DR

This work presents REPA-G, a training-free, test-time conditioning framework for diffusion models that leverages representation-aligned visual features from a frozen backbone to steer sampling toward targeted semantic concepts. Conditioning is implemented via a potential $\mathscr{V}$ and its gradient, producing a tilted distribution $\tilde{p}_{0}(x; x_c) \propto p_{0}(x) e^{\lambda \mathscr{V}(\phi(x), \phi(x_c))}$ and enabling multi-scale and multi-source composition without retraining. The authors analyze the representation space, showing semantic structuring and a well-conditioned mapping from features to densities, and demonstrate strong results on ImageNet and COCO across various IPA/SPA configurations. The approach yields controllable, diverse generations with faithful concept adherence, highlighting practical impact for flexible, perceptually precise image synthesis without textual or label-based constraints. The work also provides extensive experimental evidence and releases code for reproducibility.

Abstract

While representation alignment with self-supervised models has been shown to improve diffusion model training, its potential for enhancing inference-time conditioning remains largely unexplored. We introduce Representation-Aligned Guidance (REPA-G), a framework that leverages these aligned representations, with rich semantic properties, to enable test-time conditioning from features in generation. By optimizing a similarity objective (the potential) at inference, we steer the denoising process toward a conditioned representation extracted from a pre-trained feature extractor. Our method provides versatile control at multiple scales, ranging from fine-grained texture matching via single patches to broad semantic guidance using global image feature tokens. We further extend this to multi-concept composition, allowing for the faithful combination of distinct concepts. REPA-G operates entirely at inference time, offering a flexible and precise alternative to often ambiguous text prompts or coarse class labels. We theoretically justify how this guidance enables sampling from the potential-induced tilted distribution. Quantitative results on ImageNet and COCO demonstrate that our approach achieves high-quality, diverse generations. Code is available at https://github.com/valeoai/REPA-G.

Figures (10)

• Figure 1: Comparing class-label, text-prompt, and REPA-G conditioning. All models are trained on ImageNet. (Top) We average extracted features () from a anchor image to generate a generic "rabbit" image. (Bottom) We combine a masked feature map with a specific "lava" patch () to synthesize a rabbit on a volcano. While text prompts require lengthy descriptions and often lack precision, our feature-based conditioning offers better compositional control and provides more precise generation.
• Figure 2: Toy experiment. Comparison of the target conditional distribution sampled via the rejection method DEVROYE200683 versus our modified diffusion model. We provide additional analysis and implementation details in Appendix \ref{['apx:toy']}.
• Figure 3: Impact of representation alignment on feature space. We perform $k$-means clustering ($k=1{,}000$) on four feature spaces across ImageNet. For a reference image (with red frame), we visualize others from its assigned. Without alignment, SiT fails to form semantic groupings, making its latent space unsuitable for conditioning. In contrast, the aligned model successfully replicates the teacher's semantic structure both before and after the projection layer, resulting in semantically consistent clusters.
• Figure 4: Correlation between embedding and density distances. The narrow range of the $B/A$ ratio indicates a well-conditioned space where distances in any direction behave consistently.
• Figure 5: Semantic interpolation in the feature space. Samples are generated by bilinearly interpolating the global conditioning features between four anchor images. The smooth transitions show the semantic coherence and stability of the mapping between the embedding space and conditional densities.

Theorems & Definitions (3)

• Proposition 3.1: Proof in Apx. \ref{['sec:proof_prelim']}
• Lemma 4.3: Adapted from didi2023framework in Apx. \ref{['apx:proof_lemma']}
• Proposition 4.4: Proof in Apx. \ref{['apx:proof_lemma']}