6 papers
The Recurrent Transformer: Greater Effective Depth and Efficient Decoding
Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi +3
Transformers process tokens in parallel but are temporally shallow: at position , each layer attends to key-value pairs computed based on the previous layer, yielding a depth ca…
Opportunistic Expert Activation: Batch-Aware Expert Routing for Faster Decode Without Retraining
Costin-Andrei Oncescu, Qingyang Wu, Wai Tong Chung +5
An increasing number of LLMs employ Mixture-of-Experts (MoE) architectures where the feed-forward layer is replaced by a pool of experts and each token only activates a small subse…
Seesaw: Accelerating Training by Balancing Learning Rate and Batch Size Scheduling
Alexandru Meterez, Depen Morwani, Jingfeng Wu +3
Increasing the batch size during training -- a ''batch ramp'' -- is a promising strategy to accelerate large language model pretraining. While for SGD, doubling the batch size can…
A Simplified Analysis of SGD for Linear Regression with Weight Averaging
Alexandru Meterez, Depen Morwani, Costin-Andrei Oncescu +3
Theoretically understanding stochastic gradient descent (SGD) in overparameterized models has led to the development of several optimization algorithms that are widely used in prac…
Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond
Costin-Andrei Oncescu, Sanket Purandare, Stratos Idreos +1
While transformers have been at the core of most recent advancements in sequence generative models, their computational cost remains quadratic in sequence length. Several subquadra…
Feature emergence via margin maximization: case studies in algebraic tasks
Depen Morwani, Benjamin L. Edelman, Costin-Andrei Oncescu +2
Understanding the internal representations learned by neural networks is a cornerstone challenge in the science of machine learning. While there have been significant recent stride…