General and Efficient Steering of Diffusion Models

AI summary

Overview

Research area: Generative machine learning — controllable image synthesis with diffusion models, specifically post-hoc guidance of pretrained unconditional diffusion models toward concepts they were never trained on.

Technical level: Advanced. The paper combines diffusion sampling theory, PCA/Gaussian denoiser approximations, activation-space probing, and an iterative feature-learning method (Recursive Feature Machines), and assumes familiarity with DDPM/DDIM notation.

Scope: The paper proposes Noise-Aligned RFM Steering (NA-RFM), a post-hoc recipe that combines a high-noise PCA-based pixel correction ("noise alignment") with an intermediate/late activation-space steering direction learned by RFMs, and evaluates it on CIFAR-10, ImageNet, CelebA-HQ, and fine-grained bird species.

What This Paper Is About

Diffusion models can be guided toward a desired condition only if that condition was built into training (as in class- or text-conditioned models). Steering a pretrained model toward a condition that was never seen during training usually means either retraining it or computing classifier gradients at every denoising step, both of which are expensive. This paper introduces NA-RFM, a method that builds its guidance signals once, offline, from labeled target examples, and then samples without any inference-time gradients.

Key Contributions

1. A gradient-free post-hoc steering recipe (NA-RFM). Guidance is constructed entirely after the diffusion model is trained, from target-vs-background examples. The online sampler uses only diffusion-model forward passes, matrix-vector products, and activation-vector additions — no backpropagation through a classifier or through the diffusion model at any step.

2. A two-regime design grounded in empirical probing. Linear probes on CIFAR-10 U-Net activations (10-way probes trained on 10,000 activation examples, evaluated on a held-out 20% split) show that reverse-trajectory activations are class-informative at high noise but that forward-process activations are near chance there; forward activations become strongly class-discriminative (exceeding 80% probe accuracy) near the end of the trajectory. This motivates noise alignment at high noise and RFM activation steering at intermediate/late noise.

3. A temporal-transfer finding and its theoretical support. Cosine-similarity comparisons of forward-process activation directions across noise levels show that a direction learned at one low-noise reference level stays aligned across the intermediate window, while low-noise directions are poorly aligned with high-noise directions. A simplified model (Proposition A.1) gives conditions under which a class-discriminative direction is robust to noising.

4. Efficient sample-space computation for large activation spaces. Because the RFM average gradient outer product (AGOP) matrix is huge in activation dimension while the number of examples is small, the method recovers the needed eigenspace from the smaller N×N matrix G Gᵀ/N rather than GᵀG/N, mirroring the eigenfaces/PCA trick. PCA likewise uses the sample matrix without forming the image-space covariance.

Main Findings

• CIFAR-10 label guidance: NA-RFM reaches 96.6% target accuracy with FID 41.4, compared with 77.1% / FID 73.9 for TFG-4 and 86.0% / FID 41.9 for classifier guidance. It is also 16 times faster than TFG in the reported setting (6.2s vs. 101.7s for 16 samples with 100 DDIM steps, measured on an A100 GPU).

• Component ablations: RFM-only guidance (all steps) already yields 94.8% accuracy and FID 40.3; adding high-noise noise alignment raises accuracy to 96.6% with a small FID change to 41.4. Noise alignment alone is weaker and trades accuracy against FID: λ=3 gives 62% / FID 99, λ=8 gives 80% / FID 120, both at 5.8s.

• ImageNet 256×256: NA-RFM reaches 75.8% target accuracy with FID 98, versus 59.8% / FID 165 for TFG-4, 40.9% / 176 for TFG-1, and lower scores for DPS (38.8/193), LGD (11.5/210), FreeDoM (19.7/200), MPGD (6.8/239), and UGD (25.5/205). Evaluation uses classes 111, 222, 333, and 444 with 256 samples per class (1024 images).

• CelebA-HQ multi-attribute guidance (log-KID reported, lower is better): Gender+Age target accuracy 96.0% with log-KID −2.0, versus TFG-1's 75.2% / −3.9; TFG-4 is not reported for CelebA-HQ. Gender+Hair: 83.3% / −2.4 versus TFG-1's 76.0% / −3.6.

• Fine-grained bird species (out-of-distribution): NA-RFM reaches 14.1% accuracy with FID 72, versus 2.2% / FID 259 for TFG-4 and near-zero for the other baselines (DPS 0.0/348, LGD 0.5/246, FreeDoM 0.6/258, MPGD 0.6/249, UGD 1.1/255, TFG-1 1.3/256). The authors describe this as the largest relative gain among the three main settings where accuracy improves.

• Measured cost on ImageNet: one-time preparation of 25.45 minutes for activation collection and PCA/RFM computation, then 7.90 seconds per image at batch size 4, versus 79.70 seconds per image for TFG-4 under the same measurement. (The provided content truncates the remainder of this timing table, so the shared activation-collection and total-hours columns are not fully reported here.)

• Transfer beyond the main U-Nets: the method is reported to extend to transformer-based latent diffusion (SiT-XL/2) and to Stable Diffusion 1.5 for depth-of-field control, a photographic property the authors say can be difficult to specify reliably with text prompts alone.

• Scheduling: noise alignment is active when σ_t ≥ σ_end (high noise); RFM steering is active when σ_t ∈ [σ_R^lo, σ_R^hi] (intermediate/late), using the noise-to-signal ratio σ_t = β_t/α_t as the reporting convention. Dataset-specific ranges are given in the paper's appendix.

Methodology in Plain English

The method separates a one-time offline setup from a cheap online sampler.

Offline, for each target concept:

1. Compute PCA statistics for the target class and for the full dataset (means and leading principal directions). These are plugged into class-conditional and unconditional PCA "denoisers" — closed-form approximations of what the diffusion model would output at a given noise level.
2. Take labeled images, corrupt them to a single low-noise reference timestep, and record the activations of one chosen U-Net block (typically the last encoder block before the bottleneck). One corruption plus one denoiser pass per labeled image.
3. Train a Recursive Feature Machine on those activations with target-vs-rest labels. The RFM repeatedly fits a kernel ridge regression and reweights the feature metric using the average gradient outer product of the fitted predictor, so it progressively emphasizes activation directions that separate the target. The final steering direction is the eigenvalue-weighted, sign-corrected combination of the top k (k ∈ {1, 3, 5}) eigenvectors.

Online, during sampling:

• When noise is high, add the precomputed difference between the class-conditional and unconditional PCA denoiser outputs to the current denoised estimate, scaled by λ.
• When noise is in the intermediate/late window, run one extra denoiser pass in which the chosen block's activation tensor is shifted by the learned unit direction, scaled by the activation norm and a strength w_RFM; then extrapolate from the unsteered toward the steered denoised estimate with an amplification scale s.

Sampling uses deterministic DDIM (η = 0) with 100 steps for the main U-Net experiments, with T = 1000 training timesteps and a linear variance schedule in a VP/DDPM parameterization.

Why This Matters

The paper shows that post-hoc control of a pretrained generative model does not have to be paid for with per-step gradients, and that example data can substitute for gradient signals while giving stronger target control. For the research community this reframes guidance as a two-timescale problem — coarse structure at high noise, semantic direction at low noise — and connects diffusion steering to activation-steering ideas already common in language models.

Real-world applications:

• Personalized or niche content generation: steering a general-purpose image model toward a specific species, object class, or visual property defined only by a handful of example images.
• Photographic control: the authors demonstrate depth-of-field control with Stable Diffusion 1.5, a property that is hard to express reliably through text prompts.
• Attribute manipulation in face or portrait generation: the CelebA-HQ results target combinations such as gender plus age or gender plus hair through a post-hoc add-on rather than retraining.
• Cheap deployment of controllable synthesis: because inference avoids gradient backpropagation, guidance can be added to existing checkpoints on hardware where per-step gradient computation is impractical.

Industry relevance: any team already serving an unconditional or weakly conditioned diffusion checkpoint can add target control without retraining the model, and the reported timing gap (7.90 s/image versus 79.70 s/image on ImageNet) is directly relevant to serving costs.

Future Directions

• Automatic or adaptive scheduling of the two windows. The transition point σ_end and the RFM window [σ_R^lo, σ_R^hi] are dataset-specific and set from the paper's appendix rather than derived; a principled way to choose them for a new model or target is an open question.
• Extending beyond image classification-style targets. The current recipe assumes target-vs-rest labels; the fine-grained bird result (14.1%) suggests the hardest, most out-of-distribution targets remain far from solved.
• Scaling to larger backbones and text-conditioned models. The paper reports initial results on SiT-XL/2 and Stable Diffusion 1.5, but a systematic study of how block selection, activation-collection cost, and steering strength behave on much larger models is not reported.
• Understanding when a learned direction transfers. Proposition A.1 covers a simplified model; a fuller theory of which activation directions stay stable over which noise windows would help predict when reuse — instead of per-timestep learning — is safe.

Target Audience

Researchers and practitioners working on diffusion-model controllability, post-hoc guidance, and efficient generative inference. It is most useful to readers already comfortable with DDPM/DDIM sampling equations and classifier-free guidance, and to engineers who want to add controllable behavior to a pretrained generative model without retraining it. Readers interested in activation steering and feature-learning theory (Recursive Feature Machines) will also find the offline direction-discovery procedure relevant.

Authors’ abstract

Steering diffusion models toward conditions unseen during training typically requires either retraining with conditional inputs or per-step gradient computations, both of which incur substantial computational overhead. We present Noise-Aligned RFM Steering (NA-RFM), a general recipe for efficiently steering diffusion models without gradient guidance during inference, enabling fast controllable generation. The method combines two offline-computed signals: noise alignment, a high-noise correction from PCA statistics of the target examples and the full data, and Recursive Feature Machine (RFM) activation steering, which learns a target-discriminative direction from labeled forward-process activations. During sampling, noise alignment provides coarse control at high noise, while the RFM direction is reused over intermediate/late timesteps through lightweight activation edits. Experiments on CIFAR-10, ImageNet, CelebA, and fine-grained bird species show improved target accuracy over gradient-based post-hoc guidance baselines, improved FID on the class-guidance benchmarks, and substantial inference speedups. Code: https://github.com/isotrivial/na-rfm.