노이즈 하이퍼네트워크: 확산 모델에서 테스트 시간 계산의 분할 상환

초록

테스트 타임 스케일링(test-time scaling)이라는 새로운 패러다임은 대규모 언어 모델(LLMs, 예: 추론 모델)과 생성형 비전 모델에서 놀라운 돌파구를 마련하며, 모델이 추론 과정에서 추가적인 계산을 할당하여 점점 더 복잡한 문제를 효과적으로 해결할 수 있게 했습니다. 이러한 접근법의 개선에도 불구하고 중요한 한계가 나타납니다: 계산 시간의 상당한 증가로 인해 이 과정이 느려지고 많은 응용 분야에서 비실용적이 됩니다. 이 패러다임의 성공과 그 사용이 증가함에 따라, 우리는 추론 오버헤드를 피하면서도 그 이점을 유지하고자 합니다. 본 연구에서는 테스트 타임 스케일링 지식을 훈련 후(post-training) 모델에 통합하는 중요한 문제에 대한 하나의 해결책을 제안합니다. 구체적으로, 우리는 확산 모델(diffusion models)에서 보안 가이드 테스트 타임 노이즈 최적화를 초기 입력 노이즈를 조절하는 노이즈 하이퍼네트워크(Noise Hypernetwork)로 대체합니다. 우리는 기본 모델에 대한 충실도를 유지하면서 원하는 특성을 최적화하는 실현 가능한 노이즈 공간 목표를 통해, 증류된 생성기를 위한 보상 기울기 분포(reward-tilted distribution)를 학습하는 이론적으로 근거된 프레임워크를 제안합니다. 우리의 접근법이 명시적인 테스트 타임 최적화에서 얻는 품질 향상의 상당 부분을 훨씬 낮은 계산 비용으로 회복할 수 있음을 보여줍니다. 코드는 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

노이즈 하이퍼네트워크: 확산 모델에서 테스트 시간 계산의 분할 상환

Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models

초록

Support