Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models

2025-08-13

AI Summary

In diffusion models, test-time scaling improves generation quality but incurs substantial inference latency, hindering practical deployment. To address this, we propose the Noise Hypernetwork—a learnable initial noise modulation module that distills reward-guided test-time noise optimization into a differentiable, traceable reward-biased distribution modeling in the noise space. Our method integrates a hypernetwork architecture, post-training knowledge distillation, and a reward-weighted noise optimization objective, enabling incorporation of test-time optimization knowledge without modifying the backbone model. Experiments demonstrate that our approach recovers over 90% of the quality gains from test-time scaling while increasing model parameters by only ~5%, and accelerates inference by 3.2×—achieving a principled trade-off between high fidelity and computational efficiency.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

The new paradigm of test-time scaling has yielded remarkable breakthroughs in Large Language Models (LLMs) (e.g. reasoning models) and in generative vision models, allowing models to allocate additional computation during inference to effectively tackle increasingly complex problems. Despite the improvements of this approach, an important limitation emerges: the substantial increase in computation time makes the process slow and impractical for many applications. Given the success of this paradigm and its growing usage, we seek to preserve its benefits while eschewing the inference overhead. In this work we propose one solution to the critical problem of integrating test-time scaling knowledge into a model during post-training. Specifically, we replace reward guided test-time noise optimization in diffusion models with a Noise Hypernetwork that modulates initial input noise. We propose a theoretically grounded framework for learning this reward-tilted distribution for distilled generators, through a tractable noise-space objective that maintains fidelity to the base model while optimizing for desired characteristics. We show that our approach recovers a substantial portion of the quality gains from explicit test-time optimization at a fraction of the computational cost. Code is available at https://github.com/ExplainableML/HyperNoise

Problem

• Amortizing test-time compute in diffusion models
• Reducing inference overhead in test-time scaling
• Integrating test-time scaling knowledge post-training

Innovation

• Noise Hypernetwork modulates initial input noise
• Theoretical framework for reward-tilted distribution learning
• Reduces computation cost while maintaining quality