Learning an Efficient Optimizer via Hybrid-Policy Sub-Trajectory Balance

TL;DR

This paper introduces Lo-Hp, a two-stage weight generation framework that improves optimizer flexibility and efficiency by combining on-policy and off-policy learning, addressing long-horizon and over-coupling issues in neural network weight generation.

Contribution

It proposes a hybrid-policy sub-trajectory balance objective for decoupled weight generation, enabling better local and global optimization policies.

Findings

Achieves higher accuracy in transfer and few-shot learning tasks.

Demonstrates improved inference efficiency over existing methods.

Addresses long-horizon issues in weight optimization.

Abstract

Recent advances in generative modeling enable neural networks to generate weights without relying on gradient-based optimization. However, current methods are limited by issues of over-coupling and long-horizon. The former tightly binds weight generation with task-specific objectives, thereby limiting the flexibility of the learned optimizer. The latter leads to inefficiency and low accuracy during inference, caused by the lack of local constraints. In this paper, we propose Lo-Hp, a decoupled two-stage weight generation framework that enhances flexibility through learning various optimization policies. It adopts a hybrid-policy sub-trajectory balance objective, which integrates on-policy and off-policy learning to capture local optimization policies. Theoretically, we demonstrate that learning solely local optimization policies can address the long-horizon issue while enhancing the generation of global optimal weights. In addition, we validate Lo-Hp's superior accuracy and inference efficiency in tasks that require frequent weight updates, such as transfer learning, few-shot learning, domain generalization, and large language model adaptation.