Technical Summary: Test-Time Conditioning with Representation-Aligned Visual Features (REPA-G)

Problem Statement

While representation alignment with self-supervised models has been shown to accelerate the training of diffusion models and improve sample quality, its potential for enhancing inference-time conditioning remains largely unexplored. Current conditioning strategies in image synthesis typically rely on class labels or textual descriptions. Class-based conditioning offers only coarse control limited to predefined categories, while text prompts, though flexible, are often imprecise and ambiguous. Furthermore, existing post-training conditioning methods often require additional parameters, task-specific training, or architectural modifications that remain fixed at inference. There is a need for a method that enables precise, flexible, and fine-grained control over generation without retraining the model or relying on ambiguous prompts.

Methodology: REPA-G

The authors introduce Representation-Aligned Guidance (REPA-G), a framework that leverages the internal representations of diffusion models trained with representation alignment to enable test-time visual conditioning.

Core Mechanism

REPA-G operates by optimizing a similarity objective (a "potential") during the inference process. The method guides the denoising process toward a target representation extracted from a pre-trained self-supervised feature extractor (e.g., DINOv2).

1. Feature Extraction: Visual tokens are extracted from a reference (anchor) image using the same self-supervised network used during the model's training alignment phase.
2. Gradient Guidance: The sampling procedure is modified by computing the gradient of a potential function $\mathscr{V}$ that quantifies the alignment between the generated sample's internal features and the conditioning features.
3. Tilted Distribution: Theoretically, this guidance steers the sampling process to draw from a "tilted distribution" $\tilde{p}_0(x; x_c) \propto p_0(x) e^{\lambda \mathscr{V}(\phi(x), \phi(x_c))}$, where $p_0$ is the original data distribution and $\lambda$ controls the strength of the guidance.

Design Choices for the Potential

The framework supports multiple granularities of control by adapting the potential function $\mathscr{V}$:

• Independent Patch Alignment (IPA): Allows for spatial flexibility.
  • Full Feature Map: Enforces dense spatial constraints for stochastic reconstruction.
  • Masked Feature Map: Preserves information within specific regions (e.g., object shape) while allowing variation elsewhere.
  • Average Concept: Aligns the global semantic average, removing spatial constraints to guide the general category (e.g., "rabbit").
  • Single Concept: Focuses guidance on a specific local patch.
• Selective Patch Alignment (SPA): Uses a soft-max mechanism with a temperature parameter $T$ to dynamically weight patches based on their similarity to a target concept. This allows the model to adaptively localize concepts, bypassing rigid spatial constraints.

Prerequisites

The method requires the diffusion model to be trained with representation alignment (e.g., using REPA or REPA-E objectives). This ensures the model's internal feature space is semantically meaningful and well-conditioned, meaning distances in the feature space correlate with distances in the conditional density space. Without this alignment, the feature space lacks the semantic structure necessary for effective conditioning.

Key Contributions

1. REPA-G Framework: The first test-time visual conditioning framework for diffusion models trained with representation alignment. It aligns internal diffusion features with targets from a self-supervised encoder without requiring fine-tuning or retraining.
2. Theoretical Justification: The authors provide a theoretical analysis showing that this guidance enables sampling from a potential-induced tilted distribution. They also justify that self-supervised features are well-suited for this task due to the semantic coherence and smooth mapping properties of the aligned representation space.
3. Versatile Control: The method enables control over concrete and abstract visual concepts at multiple scales, ranging from fine-grained texture matching via single patches to broad semantic guidance using global tokens. It also supports multi-concept composition (e.g., combining an object from one image with a background from another).

Experimental Results

The authors evaluated REPA-G on ImageNet and COCO datasets using flow matching models (SiT) trained with representation alignment.

• Distribution-Level Performance: On ImageNet, REPA-G significantly outperforms unconditional baselines and standard SiT models when conditioned. For instance, using REPA-E with DINOv2 features (FDINO), the method achieved a Fréchet Inception Distance (FID) of 1.45 (Full Feature Map) and 2.30 (Masked), compared to 15.07 for the unconditional baseline. The method also demonstrated high Inception Scores (IS) and Precision/Recall metrics.
• Instance-Level Alignment: The generated images showed high cosine similarity to the anchor images across DINOv2, JEPA, and CLIP feature spaces, with PSNR values indicating strong pixel-level fidelity in masked reconstructions.
• Comparison with Text-to-Image: When compared against a text-to-image baseline (CAD-I) on COCO, REPA-G using visual features achieved superior performance (FID of 2.11 vs. 37.98 for text), suggesting visual features provide a denser and more informative signal than text captions.
• Multi-Concept Composition: In compositional generation tasks (combining anchor objects with target backgrounds), REPA-G variants (IPA and SPA) significantly outperformed class-conditional and interpolated baselines in terms of CLIP scores and patch-wise similarity metrics.

Significance and Claims

The paper claims that REPA-G offers a flexible and precise alternative to often ambiguous text prompts or coarse class labels. By operating entirely at inference time, it avoids the computational cost and complexity of training new conditioning mechanisms or fine-tuning models. The authors emphasize that the effectiveness of the method is contingent upon the model being trained with representation alignment, which ensures the internal feature space possesses the necessary semantic structure. While the approach currently requires tuning the guidance scale $\lambda$ for compositional tasks and faces challenges with varying scene viewpoints, it represents a significant step toward controllable, feature-driven generative modeling.