6 papers
RL Excursions during Pre-Training: Re-examining Policy Optimization for LLM training
Rachit Bansal, Clara Mohri, Tian Qin +2
The standard LLM training pipeline applies reinforcement learning (RL) only after pre-training and supervised fine-tuning (SFT). We question this status quo by training a LLM from…
Why Larger Models Learn More: Effects of Capacity, Interference, and Rare-Task Retention
Jing Huang, Daniel Wurgaft, Rachit Bansal +6
Larger models learn tasks smaller models do not. What drives this phenomenon? We develop a simple phenomenological argument that power-law scaling already suggests that a larger mo…
Interleaved Head Attention
Sai Surya Duvvuri, Chanakya Ekbote, Rachit Bansal +6
Multi-Head Attention (MHA) is the core computational primitive underlying modern Large Language Models (LLMs). However, MHA suffers from a fundamental linear scaling limitation: $H…
Let's (not) just put things in Context: Test-Time Training for Long-Context LLMs
Rachit Bansal, Aston Zhang, Rishabh Tiwari +8
Progress on training and architecture strategies has enabled LLMs with millions of tokens in context length. However, empirical evidence suggests that such long-context LLMs can co…
The Art of Scaling Reinforcement Learning Compute for LLMs
Devvrit Khatri, Lovish Madaan, Rishabh Tiwari +6
Reinforcement learning (RL) has become central to training large language models (LLMs), yet the field lacks predictive scaling methodologies comparable to those established for pr…
Adam or Gauss-Newton? A Comparative Study In Terms of Basis Alignment and SGD Noise
Bingbin Liu, Rachit Bansal, Depen Morwani +3
Diagonal preconditioners are computationally feasible approximate to second-order optimizers, which have shown significant promise in accelerating training of deep learning models.…