Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models

TL;DR

HyperNoise replaces per-sample test-time noise optimization in diffusion models with a small LoRA hypernetwork that learns, once at post-training time, to modulate the initial Gaussian noise so the frozen distilled generator produces reward-tilted outputs. The framework derives a tractable noise-space objective (reward maximization + an L2 penalty on the noise modification, justified by a Lipschitz bound that lets the noise-space KL be approximated by a quadratic) and proves via the Data Processing Inequality that minimizing noise-space KL bounds the output-space KL. On GenEval, applying HyperNoise to SD-Turbo, SANA-Sprint, and FLUX-Schnell recovers a substantial portion of ReNO’s quality gains at 33×–300× lower inference cost.

Key claims

The optimal initial-noise distribution that steers a fixed distilled generator $G_\theta$ to a target reward-tilted output distribution exists in closed form ( $q^*(z) \propto p(z)\,\exp(\beta r(G_\theta(z)))$ ), so post-training reduces to learning a noise-space modulator rather than fine-tuning the generator [§3, Eq. 5].
The noise-space KL divergence $D_{KL}(q_\phi\,\|\,p)$ for a residual hypernetwork $h_\phi$ (with $z' = z + h_\phi(z)$ ) admits a tractable approximation under a Lipschitz constraint: under $\mathcal{L}(h_\phi)<1$ , the log-det error is bounded and the KL collapses to an L2 penalty $\tfrac{1}{2}\,\mathbb{E}\|h_\phi(z)\|^2$ [Theorem 1, Eq. 10–11].
Minimizing the noise-space KL upper-bounds the data-space KL via the Data Processing Inequality, providing a principled regularizer even when the output-space KL is intractable for distilled models [§3.1, Eq. 12].
The hypernetwork is instantiated as a LoRA over the generator’s own architecture with the second LoRA matrix zero-initialized; this gives $h_\phi(z) \equiv 0$ at step 0 and preserves the validity of the small-Lipschitz approximation throughout early training [§3.2 “Initialization”].
On GenEval, SD-Turbo + HyperNoise reaches 0.57 mean (vs. 0.49 base, 0.63 ReNO), with parameter count growing 0.8B → 1.1B and inference time 0.2s → 0.3s vs. ReNO’s 20s; this surpasses SDXL (0.55, 2.6B params, 6.9s) [Table 1].
On SANA-Sprint, HyperNoise reaches GenEval 0.75 vs. 0.70 base and matches LLM-based prompt optimization (also 0.75) while being 300× faster; ReNO at 30s/image reaches 0.81 [Table 1].
On FLUX-Schnell (12B), HyperNoise lifts GenEval from 0.68 to 0.72 with +1B params and +0.2s inference; ReNO needs 40s to reach 0.76 [Table 1].
Reward-hacking is the primary failure of direct reward fine-tuning of distilled generators: on the controlled “redness” task, LoRA fine-tuning of SANA-Sprint maximizes redness but collapses ImageReward, while HyperNoise achieves the redness gain without ImageReward degradation [§4.1, Figures 3–4].
HyperNoise trained at one-step inference generalizes to multi-step inference of the same generator: 1-step HyperNoise on SANA-Sprint outperforms baseline 4-step SANA-Sprint on GenEval [§4.2, Table 8].

Method

The frozen distilled generator $G_\theta: z \mapsto x$ is wrapped with a learnable residual map $z' = z + h_\phi(z)$ where $h_\phi$ is a LoRA adapter on $G_\theta$ ‘s own architecture (so weights are shared in memory) plus a modified final layer that emits only the LoRA-generated perturbation rather than adding to base weights. The training loss is

\mathcal{L}(\phi) = -\beta\,\mathbb{E}_z[r(G_\theta(z + h_\phi(z)))] + \tfrac{1}{2}\,\mathbb{E}_z\|h_\phi(z)\|^2,

with no real-data samples required — only base noise samples, the frozen generator, and the reward function. The L2 term is the Lipschitz-bound approximation to the noise-space KL against the Gaussian prior (Theorem 1); it both regularizes against reward hacking and bounds the output-space KL against $G_\theta$ ‘s native distribution via DPI. Training uses ~70k prompts from Pick-a-Picv2 + T2I-Compbench + ABC-6K; evaluation is held-out GenEval [§4.2].

Results

GenEval mean (Table 1): SD-Turbo 0.49 → +HyperNoise 0.57 (vs. +ReNO 0.63 at 20s, +BoN 0.60 at 10s, +Prompt-Opt 0.59); SANA-Sprint 0.70 → +HyperNoise 0.75 (vs. +ReNO 0.81 at 30s, +BoN 0.79 at 15s, +Prompt-Opt 0.75); FLUX-Schnell 0.68 → +HyperNoise 0.72 (vs. +ReNO 0.76 at 40s) [Table 1].
Inference cost (Table 1): HyperNoise adds 0.1s on top of base inference (0.2s → 0.3s for SD-Turbo / SANA-Sprint; 0.7s → 0.9s for FLUX-Schnell). ReNO is 33× to 300× slower (e.g. SANA-Sprint: 0.2s → 30s) [Table 1].
Multi-step generalization (Table 8): HyperNoise trained only on 1-step inference yields consistent gains over base SANA-Sprint at 1, 2, and 4 steps; 1-step HyperNoise beats 4-step base SANA-Sprint on GenEval [§4.2].
Redness reward (Figures 3–4): HyperNoise traces a strictly better redness-vs-ImageReward Pareto than LoRA fine-tuning; LoRA collapses image quality at high redness, HyperNoise does not [§4.1].
Reward-hacking comparison (Appendix C.4): direct reward fine-tuning exhibits the canonical reward-hacking artifacts observed in prior work; HyperNoise does not at matched reward [§4.2 “Superiority over Direct Fine-tuning”].

Why it’s interesting

This is the cleanest framing the wiki has of “amortize test-time compute into training” for generative vision — a question that the Inference-Time Scaling page currently lists as a sixth axis (TTT-E2E for language) without a vision analog. HyperNoise’s lever is mechanically distinct from the inference-time-scaling axes filed so far: it doesn’t add NFEs (cf. Self-Refining Video Sampling‘s Predict-and-Perturb), spawn sub-agents (cf. RLM / PARL), or compress the cache (cf. Inference-Time Hyper-Scaling with KV Cache Compression); it moves the optimization step itself from per-sample inference to one-time post-training. The “frozen distilled generator + tiny LoRA + structured input modulation” recipe is the same template Parameter-Efficient Finetuning has consolidated for SDR→HDR retargeting (DiffHDR: Re-Exposing LDR Videos with Video Diffusion Models, HDR Video Generation via Latent Alignment with Logarithmic Encoding (LumiVid)) and structured-conditioning slot replacement (FIT: A Large-Scale Dataset for Fit-Aware Virtual Try-On) — HyperNoise extends it to the noise slot. It also gives Diffusion Distillation a complementary post-training stage: distill first to 1–4 NFE, then amortize reward signal into a noise hypernetwork. The reward-hacking discussion runs parallel to TDM-R1: Reinforcing Few-Step Diffusion Models with Non-Differentiable Reward‘s TDM-R1, but reaches a different fix — TDM-R1 uses a Surrogate Reward + marginal-KL anchor on the generator; HyperNoise sidesteps the generator entirely by working in noise space, where the prior is closed-form Gaussian and the KL is tractable.