Learning to Generalize Provably in Learning to Optimize

TL;DR

This paper introduces a novel approach to improve the generalization of learned optimizers and optimizees by incorporating flatness-aware regularizers based on local entropy and Hessian metrics, supported by theoretical analysis and extensive experiments.

Contribution

It proposes a new method to enhance generalization in learning to optimize by integrating flatness metrics into the meta-training process, with theoretical justification and empirical validation.

Findings

Significantly improved generalization on multiple models and optimizees.

Theoretical connection between local entropy and Hessian for landscape analysis.

Effective incorporation of flatness-aware regularizers into L2O framework.

Abstract

Learning to optimize (L2O) has gained increasing popularity, which automates the design of optimizers by data-driven approaches. However, current L2O methods often suffer from poor generalization performance in at least two folds: (i) applying the L2O-learned optimizer to unseen optimizees, in terms of lowering their loss function values (optimizer generalization, or ``generalizable learning of optimizers"); and (ii) the test performance of an optimizee (itself as a machine learning model), trained by the optimizer, in terms of the accuracy over unseen data (optimizee generalization, or ``learning to generalize"). While the optimizer generalization has been recently studied, the optimizee generalization (or learning to generalize) has not been rigorously studied in the L2O context, which is the aim of this paper. We first theoretically establish an implicit connection between the local entropy and the Hessian, and hence unify their roles in the handcrafted design of generalizable optimizers as equivalent metrics of the landscape flatness of loss functions. We then propose to incorporate these two metrics as flatness-aware regularizers into the L2O framework in order to meta-train optimizers to learn to generalize, and theoretically show that such generalization ability can be learned during the L2O meta-training process and then transformed to the optimizee loss function. Extensive experiments consistently validate the effectiveness of our proposals with substantially improved generalization on multiple sophisticated L2O models and diverse optimizees. Our code is available at: https://github.com/VITA-Group/Open-L2O/tree/main/Model_Free_L2O/L2O-Entropy.