LionVote: Per-Layer Learning Rate Adaptation for Lion

2026-07-10Machine Learning

Machine Learning
AI summary

The authors found that the usual learning rate used for training Vision Transformer (ViT) models is unevenly too high across different types of model layers. To fix this, they created a method called LionVote that adjusts the learning rate for each layer separately, based on signals like gradient stability and momentum. This layer-specific adjustment led to slightly better accuracy on a small ViT model trained on CIFAR-100 compared to standard Lion and AdamW optimizers. However, the benefits of this approach depend on the model architecture and task, with traditional methods still working best for some convolutional networks.

learning rateVision Transformer (ViT)layer-wise adaptationLion optimizerAdamWgradient stabilitymomentumCIFAR-100cosine annealingSGD
Authors
Kris Atallah
Abstract
Per-layer diagnostics reveal that, at the prescribed learning rate, Lion's effective scale is 2.6-2.8x too high for attention and MLP parameters and ~2x too high for normalization layers on ViT-Tiny/CIFAR-100; this 32% cross-layer-type disparity cannot be reproduced by a single global rate. The measurement comes from LionVote, a per-layer learning rate mechanism in which each parameter tensor maintains a compound level, a persistent integer updated every c epochs by two diagnostics (gradient direction stability and momentum health) resolved by a validation loss tiebreaker. Voting thresholds derive from geometric identities, the EMA time constant, and a noise-floor estimate; cadence is bounded structurally and selected by ablation. On ViT-Tiny/CIFAR-100, LionVote achieves 69.7% top-1 accuracy vs. Lion's 69.0% (p < 0.02, Welch's t-test) and AdamW's 68.8%. Per-layer adaptation value depends on both architectural heterogeneity and task; on uniform CNN architectures tuned SGD with cosine annealing remains dominant, and on ViT architectures gains are task-dependent.