GRAIN: Group Aggregation via Min-Norm Objective

2026-06-22

AI Summary

This work addresses the training instability and high inter-run performance variance commonly observed in deep learning when using over-parameterized models with limited downstream data. To this end, the authors propose GRAIN, a lightweight training algorithm that, for the first time, incorporates a minimum-norm objective into gradient aggregation. Instead of the conventional arithmetic averaging of mini-batch gradients, GRAIN employs a convex combination to coordinate gradients within a single backward pass—without incurring additional computational overhead. Theoretical analysis shows that GRAIN achieves a tighter uniform stability bound than standard SGD, and its solution almost surely differs from the mean-based solution. Empirical results consistently demonstrate improved average performance and substantially reduced variance across generative, classification, and regression tasks, with notable effectiveness in large-scale pre-trained model settings.

Abstract

Learning instability is a long-standing problem across machine learning, but it is especially acute in the overparameterized regime that defines modern deep learning: large models fine-tuned or trained on limited data traverse flat loss landscapes with many nearly-equivalent minima, and stochastic factors (initialization, data order, dropout, hardware non-determinism) can route optimization to very different solutions. The rise of large pretrained models (LPMs) makes the problem more urgent: training cost is high, downstream data is often small, and repeated runs for variance reduction are prohibitive. We introduce \textbf{GRAIN} (\textbf{G}roup \textbf{A}ggregation via m\textbf{IN}-norm objective), a lightweight training algorithm that replaces the mean aggregation used in mini-batch optimization (both across mini-batches and within a mini-batch) with a min-norm convex combination of group-wise gradients. \mName guarantees a non-negative inner product between the aggregated update and every group gradient, resolving intra- and inner-batch gradient conflict, and retains an $\mathcal{O}(1/T)$ convergence rate comparable to SGD. Under mild smoothness and absolute-continuity assumptions, the min-norm solution differs almost surely from the arithmetic mean, which yields a uniform-stability bound for \mName strictly tighter than the standard bound for SGD. Empirically across generation, classification, and regression at LPM scale, \mName delivers consistent improvements in mean performance and reductions in run-to-run variance over a broad suite of tasks, with no extra training-time or storage cost beyond a single backward pass.

Problem

• training instability
• overparameterized models
• large pretrained models
• gradient conflict
• run-to-run variance

Innovation

• min-norm aggregation
• gradient conflict resolution
• uniform stability
• overparameterized learning
• large pretrained models