Meta-Learning Adaptive Loss Functions

TL;DR

AdaLFL tackles the short-horizon bias of offline loss-function learning by performing online meta-learning of a loss function that is updated after every base-parameter update. The loss is parameterized as a small neural network and optimized via unrolled differentiation, allowing the loss to adapt in shape and scale throughout training. Empirical results across regression and classification tasks show AdaLFL improves convergence speed and final accuracy over handcrafted losses and offline baselines, while revealing phenomena such as implicit learning-rate scheduling, early stopping regularization, and adaptive label smoothing. This work demonstrates the practical viability and insights of online loss-function learning, suggesting broad applicability in deep learning pipelines without manual loss design.

Abstract

Loss function learning is a new meta-learning paradigm that aims to automate the essential task of designing a loss function for a machine learning model. Existing techniques for loss function learning have shown promising results, often improving a model's training dynamics and final inference performance. However, a significant limitation of these techniques is that the loss functions are meta-learned in an offline fashion, where the meta-objective only considers the very first few steps of training, which is a significantly shorter time horizon than the one typically used for training deep neural networks. This causes significant bias towards loss functions that perform well at the very start of training but perform poorly at the end of training. To address this issue we propose a new loss function learning technique for adaptively updating the loss function online after each update to the base model parameters. The experimental results show that our proposed method consistently outperforms the cross-entropy loss and offline loss function learning techniques on a diverse range of neural network architectures and datasets.

Figures (18)

• Figure 1: Example adaptive meta-learned loss functions generated by AdaLFL on the CIFAR-10 dataset, where each row represents a classification loss function, and the color represents the current gradient step.
• Figure 2: Computational graph of AdaLFL, where $\theta$ is updated using $\mathcal{M}_{\phi}$ in the inner loop (Base Update). The optimization path is tracked in the computational graph and then used to update $\phi$ based on the meta-objective in the outer loop (Meta Update). The dashed lines show the gradients for $\theta$ and $\phi$ with respect to their given objectives.
• Figure 3: Mean learning curves across 10 independent executions of each algorithm on each task + model pair, showing the log of the training mean squared error or error rate (y-axis) against gradient steps (x-axis). Best viewed in color.
• Figure 4: Loss functions generated by AdaLFL on the CIFAR-10 dataset, where each row represents a loss function and the color represents the current gradient step.
• Figure 5: Four example loss functions generated by AdaLFL using the network architecture proposed in bechtle2021meta, which uses a softplus activation in the output layer, causing flattening behavior degrading learning performance.