arXiv subjects

Aditya Somasundaram

Publications and source records attributed to Aditya Somasundaram.

3 recordsLinked to original sources

arXiv2026

Toward a First-Principles Update Geometry for the Language-Model Head

Muon motivates designing optimizer geometry around the function of each parameter block and uses the spectral norm for hidden linear layers. For the language-model head, the spectral norm is not a faithful measure of functional change. Softmax removes shared logit shifts, whereas the spectral norm can assign arbitrarily large size to updates that change no output probability. We therefore treat the LM head and softmax as one module and derive an update geometry for their composition. Hilbert's projective distance respects this invariance as it measures the largest change in pairwise log odds. For an update $S$ with token rows $s_i^\top$, we show that the largest Hilbert distance over $\left\lVert h\right\rVert_2\leq H$ is exactly $H D(S)$, where $D(S)=\max_{i<j}\left\lVert s_i - s_j\right\rVert_2$ is the Euclidean row diameter. This diameter replaces the spectral norm in the resulting Muon-style steepest descent problem. An exact solution is possible, but its direct formulation contains one $d$-dimensional vector variable for every token pair. For a vocabulary size of approximately $50$k, this means more than one billion token pairs, making the calculation impractical at every training step. We instead impose a stronger common-ball constraint and derive projected RowNorm as an $O(Vd)$ solution. For the exact RowNorm oracle, we prove that its first-order decrease is at least $1/\sqrt{2}$ of the exact diameter-constrained optimum. With Muon on the backbone, experiments across three seeds at 190M, 380M, and 640M parameters show that RowNorm reduces mean final step diameters and empirical Hilbert RMS perturbations by factors of $45$--$60$ and $12$--$15$, respectively, with only a $0.0057$--$0.0153$ increase in mean final validation loss.

cs.LG↗

arXiv2025

Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful

Conventional wisdom dictates that small batch sizes make language model pretraining and fine-tuning unstable, motivating gradient accumulation, which trades off the number of optimizer steps for a proportional increase in batch size. While it is common to decrease the learning rate for smaller batch sizes, other hyperparameters are often held fixed. In this work, we revisit small batch sizes all the way down to batch size one, and we propose a rule for scaling Adam hyperparameters to small batch sizes. In particular, rather than holding the decay rate of the second moment fixed across batch sizes, we propose to hold its half-life fixed in terms of tokens. We find that small batch sizes (1) train stably, (2) are consistently more robust to hyperparameter choices, (3) achieve equal or better per-FLOP performance than larger batch sizes, and (4) notably enable stable language model training with vanilla SGD, even without momentum, despite storing no optimizer state. Building on these results, we provide practical recommendations for selecting a batch size and setting optimizer hyperparameters. We further recommend against gradient accumulation unless training on multiple devices with multiple model replicas. Finally, we show that a small batch size combined with an optimizer with a small state size can provide the performance benefits of full fine-tuning while maintaining a similar memory footprint to LoRA.

cs.LG↗

arXiv2025

Learning Using a Single Forward Pass

We propose a learning algorithm to overcome the limitations of traditional backpropagation in resource-constrained environments: Solo Pass Embedded Learning Algorithm (SPELA). SPELA operates with local loss functions to update weights, significantly saving on resources allocated to the propagation of gradients and storing computational graphs while being sufficiently accurate. Consequently, SPELA can closely match backpropagation using less memory. Moreover, SPELA can effectively fine-tune pre-trained image recognition models for new tasks. Further, SPELA is extended with significant modifications to train CNN networks, which we evaluate on CIFAR-10, CIFAR-100, and SVHN 10 datasets, showing equivalent performance compared to backpropagation. Our results indicate that SPELA, with its features such as local learning and early exit, is a potential candidate for learning in resource-constrained edge AI applications.

cs.AI↗