Generalizable Learning to Optimize into Wide Valleys

Junjie Yang, Tianlong Chen, Mingkang Zhu, Fengxiang He, Dacheng Tao, Yingbin Liang, Zhangyang Wang

Keywords: L2O, Generalization, Flatness, Entropy-SGD

Abstract: Learning to optimize (L2O) has gained increasing popularity in various optimization tasks, since classical optimizers usually require laborious, problem-specific design and hyperparameter tuning. However, current L2O approaches are designed for fast minimization of the objective function value (i.e., training error), hence often suffering from poor generalization ability such as in training deep neural networks (DNNs), including ($i$) disappointing performance across unseen optimizees $\textit{(optimizer generalization)}$; ($ii$) unsatisfactory test-set accuracy of trained DNNs ($\textit{optmizee generalization}$). To overcome the limitations, this paper introduces $\textit{flatness-aware}$ regularizers into L2O for shaping the local geometry of optimizee's loss landscape. Specifically, it guides optimizee to locate well-generalizable minimas in large flat regions of loss surface, while tending to avoid sharp valleys. Such optimizee generalization abilities of $\textit{flatness-aware}$ regularizers have been proved theoretically. Extensive experiments consistently validate the effectiveness of our proposals with substantially improved generalization on multiple sophisticated L2O models and diverse optimizees. Our theoretical and empirical results solidify the foundation for L2O's practically usage. All codes and pre-trained models will be shared upon acceptance.

One-sentence Summary: This paper introduces flatness-aware regularizers into L2O which guide optimizee to converge into wide valleys in loss landscape.