Test-Time Conditioning with Representation-Aligned Visual Features

TL;DR

This paper introduces REPA-G, a test-time guidance framework that uses representation alignment to steer generative models towards desired features, enabling versatile and precise control during inference.

Contribution

The paper proposes a novel inference-time guidance method using representation-aligned features, enhancing control and diversity in generative models without retraining.

Findings

Achieves high-quality, diverse generations on ImageNet and COCO

Enables multi-concept composition for complex guidance

Operates entirely at inference time for flexibility

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.