GRAIN: Group Aggregation via Min-Norm Objective

Modern deep learning, particularly the fine-tuning of large pretrained models (LPMs), suffers from significant training instability. Small stochastic factors—such as random initialization, data ordering, and hardware non-determinism—can lead to vastly different model performance. The authors investigate whether this instability is driven by 'gradient cancellation,' where conflicting gradients from different data groups average out to ineffective updates, and propose a new optimization approach to mitigate this.

To address gradient conflict, the authors introduce GRAIN (Group Aggregation via mIN-norm objective). Instead of using the arithmetic mean of gradients within a mini-batch, GRAIN partitions the batch into $m$ groups and solves a convex quadratic program to find the min-norm convex combination of these group-wise gradients. This ensures that the aggregated update maintains a non-negative inner product with every group gradient, effectively preventing the cancellation of opposing signals. The method is optimizer-agnostic and can be implemented with minimal overhead, as it requires only a single backward pass and a small optimization step.

Theoretical analysis shows that GRAIN guarantees convergence to a stationary point at an $\mathcal{O}(1/T)$ rate and provides a strictly tighter uniform-stability bound compared to standard SGD. Empirically, GRAIN demonstrates superior consistency across a wide range of tasks, including generative LLM fine-tuning, sequence classification, and image classification under distribution shift. Notably, GRAIN consistently avoids the 'collapsed failures' observed in baseline methods (such as LoRA, SWA, or SAM) and significantly reduces run-to-run variance without increasing training time or storage requirements.

In the era of large pretrained models, where training is computationally expensive and downstream data is often scarce, repeated runs to average out seed-induced variance are frequently infeasible. By directly targeting the optimization process to resolve gradient conflict, GRAIN provides a reliable way to achieve stable, high-performing models in a single training run, which is critical for safety-critical or domain-specific applications.