← Search

Albert Gu

36 accepted papers

2026

Joint Distillation for Fast Likelihood Evaluation and Sampling in Flow-based Models

ICLR 2026poster

Log-likelihood evaluation enables important capabilities in generative models, including model comparison, certain fine-tuning objectives, and many downstream applications. Yet paradoxically, some of today's best generative models -- diffusion and flow-based models -- still require hundreds to thous…

Cited by 0SourceScholar
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

dnaHNet: A Scalable and Hierarchical Foundation Model for Genomic Sequence Learning

ICML 2026oral

Genomic foundation models have the potential to decode DNA syntax, yet face a fundamental tradeoff. Standard subword tokenizers fragment biologically meaningful motifs such as codons and regulatory elements, while nucleotide-level models preserve biological coherence but incur prohibitive computatio…

Cited by 0SourceScholar
2025

On the Benefits of Memory for Modeling Time-Dependent PDEs

ICLR 2025oral

Data-driven techniques have emerged as a promising alternative to traditional numerical methods for solving PDEs. For time-dependent PDEs, many approaches are Markovian---the evolution of the trained system only depends on the current state, and not the past states. In this work, we investigate the…

Cited by 3SourcePDFScholar
2025

Understanding the Skill Gap in Recurrent Language Models: The Role of the Gather-and-Aggregate Mechanism

ICML 2025poster

State-space models (SSMs) offer efficient alternatives to Transformers for long sequences, but their fixed-size recurrent state limits capability on algorithmic tasks, such as retrieving past context. In this work, we examine how in-context retrieval operates in Transformer- and SSM-based language m…

Cited by 1SourcePDFScholar
2024

Augmenting Conformers With Structured State-Space Sequence Models For Online Speech Recognition

ICASSP 2024accepted

Online speech recognition, where the model only accesses context to the left, is an important and challenging use case for ASR systems. In this work, we investigate augmenting neural encoders for online ASR by incorporating structured state-space sequence models (S4), a family of models that provide…

Cited by 0SourceScholar
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

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

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…

2024

Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models

NeurIPS 2024poster

Transformer architectures have become a dominant paradigm for domains like language modeling but suffer in many inference settings due to their quadratic-time self-attention. Recently proposed subquadratic architectures, such as Mamba, have shown promise, but have been pretrained with substantially…

Cited by 21SourcePDFScholar
2023

How to Train your HIPPO: State Space Models with Generalized Orthogonal Basis Projections

ICLR 2023poster

Linear time-invariant state space models (SSM) are a classical model from engineering and statistics, that have recently been shown to be very promising in machine learning through the Structured State Space sequence model (S4). A core component of S4 involves initializing the SSM state matrix to a…

2023

Modelling Long Range Dependencies in $N$D: From Task-Specific to a General Purpose CNN

ICLR 2023poster

Performant Convolutional Neural Network (CNN) architectures must be tailored to specific tasks in order to consider the length, resolution, and dimensionality of the input data. In this work, we tackle the need for problem-specific CNN architectures. We present the Continuous Convolutional Neural Ne…

2023

Resurrecting Recurrent Neural Networks for Long Sequences

ICML 2023oral

Recurrent Neural Networks (RNNs) offer fast inference on long sequences but are hard to optimize and slow to train. Deep state-space models (SSMs) have recently been shown to perform remarkably well on long sequence modeling tasks, and have the added benefits of fast parallelizable training and RNN-…

Cited by 294SourcePDFScholar
2023

Structured State Space Models for In-Context Reinforcement Learning

NeurIPS 2023poster

Structured state space sequence (S4) models have recently achieved state-of-the-art performance on long-range sequence modeling tasks. These models also have fast inference speeds and parallelisable training, making them potentially useful in many reinforcement learning settings. We propose a modif…

2022

It’s Raw! Audio Generation with State-Space Models

ICML 2022oral

Developing architectures suitable for modeling raw audio is a challenging problem due to the high sampling rates of audio waveforms. Standard sequence modeling approaches like RNNs and CNNs have previously been tailored to fit the demands of audio, but the resultant architectures make undesirable co…

2022

On the Parameterization and Initialization of Diagonal State Space Models

NeurIPS 2022accept

State space models (SSM) have recently been shown to be very effective as a deep learning layer as a promising alternative to sequence models such as RNNs, CNNs, or Transformers. The first version to show this potential was the S4 model, which is particularly effective on tasks involving long-rang…

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
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

HoroPCA: Hyperbolic Dimensionality Reduction via Horospherical Projections

ICML 2021spotlight

This paper studies Principal Component Analysis (PCA) for data lying in hyperbolic spaces. Given directions, PCA relies on: (1) a parameterization of subspaces spanned by these directions, (2) a method of projection onto subspaces that preserves information in these directions, and (3) an objective…

2021

Model Patching: Closing the Subgroup Performance Gap with Data Augmentation

ICLR 2021poster

Classifiers in machine learning are often brittle when deployed. Particularly concerning are models with inconsistent performance on specific subgroups of a class, e.g., exhibiting disparities in skin cancer classification in the presence or absence of a spurious bandage. To mitigate these performan…

2020

From Trees to Continuous Embeddings and Back: Hyperbolic Hierarchical Clustering

NeurIPS 2020poster

Similarity-based Hierarchical Clustering (HC) is a classical unsupervised machine learning algorithm that has traditionally been solved with heuristic algorithms like Average-Linkage. Recently, Dasgupta reframed HC as a discrete optimization problem by introducing a global cost function measuring th…

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

Improving the Gating Mechanism of Recurrent Neural Networks

ICML 2020poster

Gating mechanisms are widely used in neural network models, where they allow gradients to backpropagate easily through depth or time. However, their saturation property introduces problems of its own. For example, in recurrent models these gates need to have outputs near 1 to propagate information o…

Cited by 108SourcePDFScholar
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
2020

No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems

NeurIPS 2020poster

In real-world classification tasks, each class often comprises multiple finer-grained "subclasses." As the subclass labels are frequently unavailable, models trained using only the coarser-grained class labels often exhibit highly variable performance across different subclasses. This phenomenon, kn…

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

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

Learning Mixed-Curvature Representations in Product Spaces

ICLR 2019poster

The quality of the representations achieved by embeddings is determined by how well the geometry of the embedding space matches the structure of the data. Euclidean space has been the workhorse for embeddings; recently hyperbolic and spherical spaces have gained popularity due to their ability to be…

Cited by 270SourcePDFScholar
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…