40 citations · 106 across the 28 of their papers we have counts for
29 papers
TailSFT: Filtered Fine-Tuning Improves Post-Training Performance
Sadhika Malladi, Samy Jelassi, Dylan Foster +2
Reinforcement learning post-training drives reasoning and agentic capabilities in modern AI systems, yet a growing body of work shows that it is most effective when used to fine-tu…
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…
Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models
Samy Jelassi, Mujin Kwun, Rosie Zhao +5
Cross-entropy (CE) training provides dense and scalable supervision for language models, but it optimizes next-token prediction under teacher forcing rather than sequence-level beh…
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…
Let Me Think! A Long Chain-of-Thought Can Be Worth Exponentially Many Short Ones
Parsa Mirtaheri, Ezra Edelman, Samy Jelassi +2
Inference-time computation has emerged as a promising scaling axis for improving large language model reasoning. However, despite yielding impressive performance, the optimal alloc…
Echo Chamber: RL Post-training Amplifies Behaviors Learned in Pretraining
Rosie Zhao, Alexandru Meterez, Sham Kakade +3
Reinforcement learning (RL)-based fine-tuning has become a crucial step in post-training language models for advanced mathematical reasoning and coding. Following the success of fr…
To Backtrack or Not to Backtrack: When Sequential Search Limits Model Reasoning
Tian Qin, David Alvarez-Melis, Samy Jelassi +1
Recent advancements in large language models (LLMs) have significantly improved their reasoning abilities, particularly through techniques involving search and backtracking. Backtr…
The Role of Sparsity for Length Generalization in Transformers
Noah Golowich, Samy Jelassi, David Brandfonbrener +2
Training large language models to predict beyond their training context lengths has drawn much attention in recent years, yet the principles driving such behavior of length general…
LoRA Soups: Merging LoRAs for Practical Skill Composition Tasks
Akshara Prabhakar, Yuanzhi Li, Karthik Narasimhan +3
Low-Rank Adaptation (LoRA) is a popular technique for parameter-efficient fine-tuning of Large Language Models (LLMs). We study how different LoRA modules can be merged to achieve…
Collective Model Intelligence Requires Compatible Specialization
Jyothish Pari, Samy Jelassi, Pulkit Agrawal
In this work, we explore the limitations of combining models by averaging intermediate features, referred to as model merging, and propose a new direction for achieving collective…
Mixture of Parrots: Experts improve memorization more than reasoning
Samy Jelassi, Clara Mohri, David Brandfonbrener +7
The Mixture-of-Experts (MoE) architecture enables a significant increase in the total number of model parameters with minimal computational overhead. However, it is not clear what…
Universal Length Generalization with Turing Programs
Kaiying Hou, David Brandfonbrener, Sham Kakade +2
Length generalization refers to the ability to extrapolate from short training sequences to long test sequences and is a challenge for current large language models. While prior wo…
How Does Overparameterization Affect Features?
Ahmet Cagri Duzgun, Samy Jelassi, Yuanzhi Li
Overparameterization, the condition where models have more parameters than necessary to fit their training loss, is a crucial factor for the success of deep learning. However, the…
Q-Probe: A Lightweight Approach to Reward Maximization for Language Models
Kenneth Li, Samy Jelassi, Hugh Zhang +3
We present an approach called Q-probing to adapt a pre-trained language model to maximize a task-specific reward function. At a high level, Q-probing sits between heavier approache…
Repeat After Me: Transformers are Better than State Space Models at Copying
Samy Jelassi, David Brandfonbrener, Sham M. Kakade +1
Transformers are the dominant architecture for sequence modeling, but there is growing interest in models that use a fixed-size latent state that does not depend on the sequence le…
Length Generalization in Arithmetic Transformers
Samy Jelassi, Stéphane d'Ascoli, Carles Domingo-Enrich +3
We examine how transformers cope with two challenges: learning basic integer arithmetic, and generalizing to longer sequences than seen during training. We find that relative posit…
Depth Dependence of P Learning Rates in ReLU MLPs
Samy Jelassi, Boris Hanin, Ziwei Ji +3
In this short note we consider random fully connected ReLU networks of width and depth equipped with a mean-field weight initialization. Our purpose is to study the depende…
Vision Transformers provably learn spatial structure
Samy Jelassi, Michael E. Sander, Yuanzhi Li
Vision Transformers (ViTs) have achieved comparable or superior performance than Convolutional Neural Networks (CNNs) in computer vision. This empirical breakthrough is even more r…
Dissecting adaptive methods in GANs
Samy Jelassi, David Dobre, Arthur Mensch +2
Adaptive methods are a crucial component widely used for training generative adversarial networks (GANs). While there has been some work to pinpoint the "marginal value of adaptive…
Towards understanding how momentum improves generalization in deep learning
Samy Jelassi, Yuanzhi Li
Stochastic gradient descent (SGD) with momentum is widely used for training modern deep learning architectures. While it is well-understood that using momentum can lead to faster c…
Depth separation beyond radial functions
Luca Venturi, Samy Jelassi, Tristan Ozuch +1
High-dimensional depth separation results for neural networks show that certain functions can be efficiently approximated by two-hidden-layer networks but not by one-hidden-layer o…
Adaptivity without Compromise: A Momentumized, Adaptive, Dual Averaged Gradient Method for Stochastic Optimization
Aaron Defazio, Samy Jelassi
We introduce MADGRAD, a novel optimization method in the family of AdaGrad adaptive gradient methods. MADGRAD shows excellent performance on deep learning optimization problems fro…
Dual Averaging is Surprisingly Effective for Deep Learning Optimization
Samy Jelassi, Aaron Defazio
First-order stochastic optimization methods are currently the most widely used class of methods for training deep neural networks. However, the choice of the optimizer has become a…
A Permutation-Equivariant Neural Network Architecture For Auction Design
Jad Rahme, Samy Jelassi, Joan Bruna +1
Designing an incentive compatible auction that maximizes expected revenue is a central problem in Auction Design. Theoretical approaches to the problem have hit some limits in the…