Шумовые гиперсети: амортизация вычислительных затрат во время тестирования в диффузионных моделях

Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models

August 13, 2025

Авторы: Luca Eyring, Shyamgopal Karthik, Alexey Dosovitskiy, Nataniel Ruiz, Zeynep Akata

cs.AI

Аннотация

Новая парадигма масштабирования во время тестирования привела к значительным прорывам в крупных языковых моделях (LLM) (например, моделях рассуждений) и в генеративных моделях для обработки изображений, позволяя моделям выделять дополнительные вычислительные ресурсы на этапе вывода для эффективного решения всё более сложных задач. Несмотря на улучшения, достигаемые этим подходом, возникает важное ограничение: существенное увеличение времени вычислений делает процесс медленным и непрактичным для многих приложений. Учитывая успех этой парадигмы и её растущее использование, мы стремимся сохранить её преимущества, избегая при этом накладных расходов на этапе вывода. В данной работе мы предлагаем одно из решений критической проблемы интеграции знаний о масштабировании во время тестирования в модель на этапе пост-обучения. В частности, мы заменяем оптимизацию шума с управлением наградой во время тестирования в диффузионных моделях на Шумовой Гиперсеть, которая модулирует начальный входной шум. Мы предлагаем теоретически обоснованную структуру для изучения этого распределения, смещённого в сторону награды, для дистиллированных генераторов, через управляемую цель в пространстве шума, которая сохраняет верность базовой модели, одновременно оптимизируя желаемые характеристики. Мы показываем, что наш подход восстанавливает значительную часть улучшений качества, достигаемых за счёт явной оптимизации во время тестирования, при существенно меньших вычислительных затратах. Код доступен по адресу https://github.com/ExplainableML/HyperNoise.

English

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