Understanding and Improving Grokking through Cross-Sample Gradient Interactions

Grokking exposes a long delay between fitting training data and generalizing. Gradient-norm regularization can alter this delay, but the norm of the averaged gradient does not distinguish small per-sample gradients from cancellation between them. We decompose the squared batch-gradient norm into individual squared norms and positive and negative pairwise interactions. Negative interactions correspond to first-order increases in one sample’s loss under another sample’s update, yet enter norm penalties with a minus sign. We introduce Gradient Conflict-Aware Minimization (GCAM), which penalizes a coordinate-wise upper bound on these interactions using per-coordinate sign aggregation. In modular-addition Transformers, GCAM applied to attention value and output projections accelerates generalization, enables runs that fail under baseline training to generalize within the same budget, and nearly eliminates the memorization–generalization delay at narrow width. Beyond grokking, GCAM improves a vision Transformer’s CIFAR-10 accuracy without augmentation from 66.6% to 72.0% under the evaluated learning-rate choices.