← Search

Tri Dao

45 accepted papers

2026

Mamba-3: Improved Sequence Modeling using State Space Principles

ICLR 2026oral

The recent scaling of test-time compute for LLMs has restricted the practical deployment of models to those with strong capabilities that can generate high-quality outputs in an inference-efficient manner. While current Transformer-based models are the standard, their quadratic compute and linear me…

Cited by 0SourcecodeScholar
2026

Opportunistic Expert Activation: Batch-Aware Expert Routing for Faster Decode Without Retraining

ICML 2026poster

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 subset of them. During autoregressive generation, these models often enter a memory-bound regime even for moderate batch size…

Cited by 0SourceScholar
2025

Ladder-Residual: Parallelism-Aware Architecture for Accelerating Large Model Inference with Communication Overlapping

ICML 2025poster

Large language model inference is both memory-intensive and time-consuming, often requiring distributed algorithms to efficiently scale. Various model parallelism strategies are used in multi-gpu training and inference to partition computation across multiple devices, reducing memory load and comput…

2025

Long-Context State-Space Video World Models

ICCV 2025poster

Video diffusion models have recently shown promise for world modeling through autoregressive frame prediction conditioned on actions. However, they struggle to maintain long-term memory due to the high computational cost associated with processing extended sequences in attention layers. To overcome…

Cited by 0SourcePDFScholar
2024

BitDelta: Your Fine-Tune May Only Be Worth One Bit

NeurIPS 2024poster

Large Language Models (LLMs) are typically trained in two phases: pre-training on large internet-scale datasets, and fine-tuning for downstream tasks. Given the higher computational demand of pre-training, it is intuitive to assume that fine-tuning adds less new information to the model, and is thus…

2024

Caduceus: Bi-Directional Equivariant Long-Range DNA Sequence Modeling

ICML 2024poster

Large-scale sequence modeling has sparked rapid advances that now extend into biology and genomics. However, modeling genomic sequences introduces challenges such as the need to model long-range token interactions, the effects of upstream and downstream regions of the genome, and the reverse complem…

2024

FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

NeurIPS 2024spotlight

Attention, as a core layer of the ubiquitous Transformer architecture, is the bottleneck for large language models and long-context applications. elaborated an approach to speed up attention on GPUs through minimizing memory reads/writes. However, it has yet to take advantage of new capabilities pre…

2024

Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers

NeurIPS 2024poster

A wide array of sequence models are built on a framework modeled after Transformers, comprising alternating sequence mixer and channel mixer layers. This paper studies a unifying *matrix mixer* view of sequence mixers that can be conceptualized as a linear map on the input sequence. This framework…

2024

Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

ICML 2024poster

Large Language Models (LLMs) employ auto-regressive decoding that requires sequential computation, with each step reliant on the previous one's output. This creates a bottleneck as each step necessitates moving the full model parameters from High-Bandwidth Memory (HBM) to the accelerator's cache. Wh…

2024

RedPajama: an Open Dataset for Training Large Language Models

NeurIPS 2024spotlight

Large language models are increasingly becoming a cornerstone technology in artificial intelligence, the sciences, and society as a whole, yet the optimal strategies for dataset composition and filtering remain largely elusive. Many of the top-performing models lack transparency in their dataset cur…

2024

The Mamba in the Llama: Distilling and Accelerating Hybrid Models

NeurIPS 2024poster

Linear RNN architectures, like Mamba, can be competitive with Transformer models in language modeling while having advantageous deployment characteristics. Given the focus on training large-scale Transformer models, we consider the challenge of converting these pretrained models for deployment. We…

2024

Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

ICML 2024poster

While Transformers have been the main architecture behind deep learning's success in language modeling, state-space models (SSMs) such as Mamba have recently been shown to match or outperform Transformers at small to medium scale. We show that these families of models are actually quite closely rela…

2023

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

ICML 2023oral

Large language models (LLMs) with hundreds of billions of parameters have sparked a new wave of exciting AI applications. However, they are computationally expensive at inference time. Sparsity is a natural approach to reduce this cost, but existing methods either require costly retraining, have to…

2023

Effectively Modeling Time Series with Simple Discrete State Spaces

ICLR 2023poster

Time series modeling is a well-established problem, which often requires that methods (1) expressively represent complicated dependencies, (2) forecast long horizons, and (3) efficiently train over long sequences. State-space models (SSMs) are classical models for time series, and prior works combin…

2023

Hungry Hungry Hippos: Towards Language Modeling with State Space Models

ICLR 2023top-25%

State space models (SSMs) have demonstrated state-of-the-art sequence modeling performance in some modalities, but underperform attention in language modeling. Moreover, despite scaling nearly linearly in sequence length instead of quadratically, SSMs are still slower than Transformers due to poor h…

2023

Hyena Hierarchy: Towards Larger Convolutional Language Models

ICML 2023oral

Recent advances in deep learning have relied heavily on the use of large Transformers due to their ability to learn at scale. However, the core building block of Transformers, the attention operator, exhibits quadratic cost in sequence length, limiting the amount of context accessible. Existing subq…

2023

Simple Hardware-Efficient Long Convolutions for Sequence Modeling

ICML 2023poster

State space models (SSMs) have high performance on long sequence modeling but require sophisticated initialization techniques and specialized implementations for high quality and runtime performance. We study whether a simple alternative can match SSMs in performance and efficiency: directly learnin…

2022

ButterflyFlow: Building Invertible Layers with Butterfly Matrices

ICML 2022spotlight

Normalizing flows model complex probability distributions using maps obtained by composing invertible layers. Special linear layers such as masked and 1{\texttimes}1 convolutions play a key role in existing architectures because they increase expressive power while having tractable Jacobians and inv…

2022

Decentralized Training of Foundation Models in Heterogeneous Environments

NeurIPS 2022accept

Training foundation models, such as GPT-3 and PaLM, can be extremely expensive, often involving tens of thousands of GPUs running continuously for months. These models are typically trained in specialized clusters featuring fast, homogeneous interconnects and using carefully designed software system…

2022

Fine-tuning Language Models over Slow Networks using Activation Quantization with Guarantees

NeurIPS 2022accept

Communication compression is a crucial technique for modern distributed learning systems to alleviate their communication bottlenecks over slower networks. Despite recent intensive studies of gradient compression for data parallel-style training, compressing the activations for models trained with p…

Cited by 9SourcePDFScholar
2022

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

NeurIPS 2022accept

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not ach…

2022

Monarch: Expressive Structured Matrices for Efficient and Accurate Training

ICML 2022oral

Large neural networks excel in many domains, but they are expensive to train and fine-tune. A popular approach to reduce their compute or memory requirements is to replace dense weight matrices with structured ones (e.g., sparse, low-rank, Fourier transform). These methods have not seen widespread a…

2022

Pixelated Butterfly: Simple and Efficient Sparse training for Neural Network Models

ICLR 2022spotlight

Overparameterized neural networks generalize well but are expensive to train. Ideally one would like to reduce their computational cost while retaining their generalization benefits. Sparse model training is a simple and promising approach to achieve this, but there remain challenges as existing met…

2022

S4ND: Modeling Images and Videos as Multidimensional Signals with State Spaces

NeurIPS 2022accept

Visual data such as images and videos are typically modeled as discretizations of inherently continuous, multidimensional signals. Existing continuous-signal models attempt to exploit this fact by modeling the underlying signals of visual (e.g., image) data directly. However, these models have not…

Cited by 228SourcePDFScholar
2022

Transform Once: Efficient Operator Learning in Frequency Domain

NeurIPS 2022accept

Spectral analysis provides one of the most effective paradigms for information-preserving dimensionality reduction, as simple descriptions of naturally occurring signals are often obtained via few terms of periodic basis functions. In this work, we study deep neural networks designed to harness the…

2021

Catformer: Designing Stable Transformers via Sensitivity Analysis

ICML 2021spotlight

Transformer architectures are widely used, but training them is non-trivial, requiring custom learning rate schedules, scaling terms, residual connections, careful placement of submodules such as normalization, and so on. In this paper, we improve upon recent analysis of Transformers and formalize a…

Cited by 22SourcePDFScholar
2021

Combining Recurrent, Convolutional, and Continuous-time Models with Linear State Space Layers

NeurIPS 2021poster

Recurrent neural networks (RNNs), temporal convolutions, and neural differential equations (NDEs) are popular families of deep learning models for time-series data, each with unique strengths and tradeoffs in modeling power and computational efficiency. We introduce a simple sequence model inspired…

2021

Knowledge Distillation as Semiparametric Inference

ICLR 2021poster

A popular approach to model compression is to train an inexpensive student model to mimic the class probabilities of a highly accurate but cumbersome teacher model. Surprisingly, this two-step knowledge distillation process often leads to higher accuracy than training the student directly on labeled…

2021

MONGOOSE: A Learnable LSH Framework for Efficient Neural Network Training

ICLR 2021oral

Recent advances by practitioners in the deep learning community have breathed new life into Locality Sensitive Hashing (LSH), using it to reduce memory and time bottlenecks in neural network (NN) training. However, while LSH has sub-linear guarantees for approximate near-neighbor search in theory, i…

Cited by 88SourcePDFScholar
2021

Rethinking Neural Operations for Diverse Tasks

NeurIPS 2021poster

An important goal of AutoML is to automate-away the design of neural networks on new tasks in under-explored domains. Motivated by this goal, we study the problem of enabling users to discover the right neural operations given data from their specific domain. We introduce a search space of operation…

Cited by 32SourcePDFScholar
2021

Scatterbrain: Unifying Sparse and Low-rank Attention

NeurIPS 2021poster

Recent advances in efficient Transformers have exploited either the sparsity or low-rank properties of attention matrices to reduce the computational and memory bottlenecks of modeling long sequences. However, it is still challenging to balance the trade-off between model quality and efficiency to p…

2020

HiPPO: Recurrent Memory with Optimal Polynomial Projections

NeurIPS 2020spotlight

A central problem in learning from sequential data is representing cumulative history in an incremental fashion as more data is processed. We introduce a general framework (HiPPO) for the online compression of continuous signals and discrete time series by projection onto polynomial bases. Given a m…

2020

Kaleidoscope: An Efficient, Learnable Representation For All Structured Linear Maps

ICLR 2020spotlight

Modern neural network architectures use structured linear transformations, such as low-rank matrices, sparse matrices, permutations, and the Fourier transform, to improve inference speed and reduce memory usage compared to general linear maps. However, choosing which of the myriad structured transfo…

Cited by 71SourceScholar
2019

A Kernel Theory of Modern Data Augmentation

ICML 2019oral

Data augmentation, a technique in which a training set is expanded with class-preserving transformations, is ubiquitous in modern machine learning pipelines. In this paper, we seek to establish a theoretical framework for understanding data augmentation. We approach this from two directions: First,…

2019

Adaptive Hashing for Model Counting

UAI 2019poster

Randomized hashing algorithms have seen recent success in providing bounds on the model count of a propositional formula. These methods repeatedly check the satisfiability of a formula subject to increasingly stringent random constraints. Key to these approaches is the choice of a fixed family of…

2019

Approximating the Permanent by Sampling from Adaptive Partitions

NeurIPS 2019poster

Computing the permanent of a non-negative matrix is a core problem with practical applications ranging from target tracking to statistical thermodynamics. However, this problem is also #P-complete, which leaves little hope for finding an exact solution that can be computed efficiently. While the pr…

Cited by 11SourcePDFScholar
2019

Learning Fast Algorithms for Linear Transforms Using Butterfly Factorizations

ICML 2019oral

Fast linear transforms are ubiquitous in machine learning, including the discrete Fourier transform, discrete cosine transform, and other structured transformations such as convolutions. All of these transforms can be represented by dense matrix-vector multiplication, yet each has a specialized and…

2019

Low-Precision Random Fourier Features for Memory-constrained Kernel Approximation

AISTATS 2019poster

We investigate how to train kernel approximation methods that generalize well under a memory budget. Building on recent theoretical work, we define a measure of kernel approximation error which we find to be more predictive of the empirical generalization performance of kernel approximation methods…

2019

On the Downstream Performance of Compressed Word Embeddings

NeurIPS 2019spotlight

Compressing word embeddings is important for deploying NLP models in memory-constrained settings. However, understanding what makes compressed embeddings perform well on downstream tasks is challenging---existing measures of compression quality often fail to distinguish between embeddings that perfo…

2018

Learning Compressed Transforms with Low Displacement Rank

NeurIPS 2018poster

The low displacement rank (LDR) framework for structured matrices represents a matrix through two displacement operators and a low-rank residual. Existing use of LDR matrices in deep learning has applied fixed displacement operators encoding forms of shift invariance akin to convolutions. We introdu…