← Search

Beidi Chen

54 accepted papers

2026

FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernel

ICLR 2026poster

Recent advance in sparse attention mechanisms has demonstrated strong potential for reducing the computational cost of long-context training and inference in large language models (LLMs). Native Sparse Attention (NSA), one state-of-the-art approach, introduces natively trainable, hardware-aligned sp…

Cited by 0SourcecodeScholar
2026

Jackpot: Align Actor-Policy Distribution for scalable and stable RL for LLM

ICLR 2026poster

Reinforcement learning (RL) has become an increasingly important paradigm for improving large language models (LLMs) on alignment, reasoning, and coding tasks, yet it remains extremely costly. The majority of training time is spent on rollouts. Allowing actor and policy distributions to differ could…

Cited by 0SourcecodeScholar
2026

Prosperity before Collapse: How Far Can Off-Policy RL Reach with Stale Data on LLMs?

ICLR 2026poster

Reinforcement learning has been central to recent advances in large language model reasoning, but most algorithms rely on on-policy training that demands fresh rollouts at every update, limiting efficiency and scalability. Asynchronous RL systems alleviate this by decoupling rollout generation from…

Cited by 0SourcecodeScholar
2026

STEM: SCALING TRANSFORMERS WITH EMBEDDING MODULES

ICLR 2026poster

Fine-grained sparsity promises higher parametric capacity without proportional per-token compute, but often suffers from training instability, load balancing, and communication overhead. We introduce \textbf{STEM} (\emph{Scaling Transformers with Embedding Modules}), a static, token-indexed approach…

Cited by 0SourcecodeScholar
2026

ThunderAgent: A Fast, Simple, and Program-Aware Agentic Inference System

ICML 2026spotlight

Large language models (LLMs) are now used to power complex multi-turn agentic workflows. Existing services run agentic inference by assembling isolated components: an LLM inference engine (e.g., vLLM) and a tool orchestrator (e.g., Kubernetes). Although agentic workflows involve multiple LLM and too…

Cited by 0SourceScholar
2025

APE: Faster and Longer Context-Augmented Generation via Adaptive Parallel Encoding

ICLR 2025poster

Context-augmented generation (CAG) techniques, including RAG and ICL, require the efficient combination of multiple contexts to generate responses to user queries. Directly inputting these contexts as a sequence introduces a considerable computational burden by re-encoding the combined selection of…

2025

Act Only When It Pays: Efficient Reinforcement Learning for LLM Reasoning via Selective Rollouts

NeurIPS 2025poster

Reinforcement learning, such as PPO and GRPO, has powered recent breakthroughs in LLM reasoning. Scaling rollout to sample more prompts enables models to selectively use higher-quality data for training, which can stabilize RL training and improve model performance, but at the cost of significant co…

Cited by 0SourcecodeScholar
2025

GSM-$\infty$: How Do your LLMs Behave over Infinitely Increasing Reasoning Complexity and Context Length?

ICML 2025poster

Recently, long-context large language models (LLMs) have shown strong performance in information retrieval and long-document QA. However, to tackle the most challenging intellectual problems, LLMs must reason effectively in long and complex contexts (e.g., frontier mathematical research). Studying h…

Cited by 0SourcePDFScholar
2025

MagicDec: Breaking the Latency-Throughput Tradeoff for Long Context Generation with Speculative Decoding

ICLR 2025poster

Large Language Models (LLMs) have become more prevalent in long-context applications such as interactive chatbots, document analysis, and agent workflows, but it is challenging to serve long-context requests with low latency and high throughput. Speculative decoding (SD) is a widely used technique t…

2025

MagicPIG: LSH Sampling for Efficient LLM Generation

ICLR 2025spotlight

Large language models (LLMs) with long context windows have gained significant attention. However, the KV cache, stored to avoid re-computation, becomes a bottleneck. Various dynamic sparse or TopK-based attention approximation methods have been proposed to leverage the common insight that attention…

2025

Multiverse: Your Language Models Secretly Decide How to Parallelize and Merge Generation

NeurIPS 2025spotlight

Autoregressive Large Language Models (AR-LLMs) frequently exhibit implicit parallelism in sequential generation. Inspired by this, we introduce Multiverse, a new generative model enabling natively parallel generation. Multiverse internalizes a MapReduce paradigm, generating automatically through thr…

Cited by 0SourceScholar
2025

ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference

ICML 2025spotlight

With the widespread deployment of long-context large language models (LLMs), there has been a growing demand for efficient support of high-throughput inference. However, as the key-value (KV) cache expands with the sequence length, the increasing memory footprint and the need to access it for decodi…

2025

Speculative Prefill: Turbocharging TTFT with Lightweight and Training-Free Token Importance Estimation

ICML 2025poster

Improving time-to-first-token (TTFT) is an essentially important objective in modern large language model (LLM) inference engines. Optimizing TTFT directly results in higher maximal QPS and meets the requirements of many critical applications. However, boosting TTFT is notoriously challenging since…

2025

Zeroth-Order Fine-Tuning of LLMs with Transferable Static Sparsity

ICLR 2025poster

Zeroth-order optimization (ZO) is a memory-efficient strategy for fine-tuning Large Language Models using only forward passes. However, applying ZO fine-tuning in memory-constrained settings such as mobile phones and laptops remains challenging since these settings often involve weight quantization,…

Cited by 0SourcePDFScholar
2024

$\texttt{Model-GLUE}$: Democratized LLM Scaling for A Large Model Zoo in the Wild

NeurIPS 2024poster

As Large Language Models (LLMs) excel across tasks and specialized domains, scaling LLMs based on existing models has gained significant attention, which is challenged by potential performance drop when combining disparate models. Various techniques have been proposed to aggregate pre-trained LLMs,…

2024

Efficient Streaming Language Models with Attention Sinks

ICLR 2024poster

Deploying Large Language Models (LLMs) in streaming applications such as multi-round dialogue, where long interactions are expected, is urgently needed but poses two major challenges. Firstly, during the decoding stage, caching previous tokens' Key and Value states (KV) consumes extensive memory. Se…

2024

Found in the Middle: How Language Models Use Long Contexts Better via Plug-and-Play Positional Encoding

NeurIPS 2024poster

This paper aims to overcome the ``lost-in-the-middle'' challenge of large language models (LLMs). While recent advancements have successfully enabled LLMs to perform stable language modeling with up to 4 million tokens, the persistent difficulty faced by most LLMs in identifying relevant information…

2024

GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

ICML 2024oral

Training Large Language Models (LLMs) presents significant memory challenges, predominantly due to the growing size of weights and optimizer states. Common memory-reduction approaches, such as low-rank adaptation (LoRA), add a trainable low-rank matrix to the frozen pre-trained weight in each layer,…

2024

Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference

ICML 2024poster

Many computational factors limit broader deployment of large language models. In this paper, we focus on a memory bottleneck imposed by the key-value (KV) cache, a computational shortcut that requires storing previous KV pairs during decoding. While existing KV cache methods approach this problem by…

2024

HexGen: Generative Inference of Large Language Model over Heterogeneous Environment

ICML 2024poster

Serving generative inference of the large language model is a crucial component of contemporary AI applications. In this paper, our focus lies in deploying such services in a heterogeneous and cross-datacenter setting to mitigate the substantial inference costs typically associated with a single cen…

2024

JoMA: Demystifying Multilayer Transformers via Joint Dynamics of MLP and Attention

ICLR 2024poster

We propose Joint MLP/Attention (JoMA) dynamics, a novel mathematical framework to understand the training procedure of multilayer Transformer architectures. This is achieved by integrating out the self-attention layer in Transformers, producing a modified dynamics of MLP layers only. JoMA removes un…

2024

KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache

ICML 2024poster

Efficiently serving large language models (LLMs) requires batching many requests together to reduce the cost per request. Yet, the key-value (KV) cache, which stores attention keys and values to avoid re-computations, significantly increases memory demands and becomes the new bottleneck in speed and…

2024

LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

ACL 2024long

We present LayerSkip, an end-to-end solution to speed-up inference of large language models (LLMs). First, during training we apply layer dropout, with low dropout rates for earlier layers and higher dropout rates for later layers, and an early exit loss where all transformer layers share the same e…

2024

Learn To be Efficient: Build Structured Sparsity in Large Language Models

NeurIPS 2024spotlight

Large Language Models (LLMs) have achieved remarkable success with their billion-level parameters, yet they incur high inference overheads. The emergence of activation sparsity in LLMs provides a natural approach to reduce this cost by involving only parts of the parameters for inference. However, e…

Cited by 14SourcePDFScholar
2024

LoCoCo: Dropping In Convolutions for Long Context Compression

ICML 2024poster

This paper tackles the memory hurdle of of processing long context sequences in Large Language Models (LLMs), by presenting a novel approach, Dropping In Convolutions for **Lo**ng **Co**ntext **Co**mpression (**LoCoCo**). LoCoCo employs only a fixed-size Key-Value (KV) cache, and can enhance efficie…

2024

Megalodon: Efficient LLM Pretraining and Inference with Unlimited Context Length

NeurIPS 2024poster

The quadratic complexity and weak length extrapolation of Transformers limits their ability to scale to long sequences, and while sub-quadratic solutions like linear attention and state space models exist, they empirically underperform Transformers in pretraining efficiency and downstream task accur…

2024

Mini-Sequence Transformers: Optimizing Intermediate Memory for Long Sequences Training

NeurIPS 2024poster

We introduce Mini-Sequence Transformer (MsT), a simple and effective methodology for highly efficient and accurate LLM training with extremely long sequences. MsT partitions input sequences and iteratively processes mini-sequences to reduce intermediate memory usage. Integrated with activation recom…

Cited by 0SourcePDFScholar
2024

Nearest Neighbor Speculative Decoding for LLM Generation and Attribution

NeurIPS 2024poster

Large language models (LLMs) often hallucinate and lack the ability to provide attribution for their generations. Semi-parametric LMs, such as kNN-LM, approach these limitations by refining the output of an LM for a given prompt using its nearest neighbor matches in a non-parametric data store. Howe…

2024

On the Surprising Effectiveness of Attention Transfer for Vision Transformers

NeurIPS 2024poster

Conventional wisdom suggests that pre-training Vision Transformers (ViT) improves downstream performance by learning useful representations. Is this actually true? We investigate this question and find that the features and representations learned during pre-training are not essential. Surprisingly…

2024

S$^{2}$FT: Efficient, Scalable and Generalizable LLM Fine-tuning by Structured Sparsity

NeurIPS 2024poster

Current PEFT methods for LLMs can achieve high quality, efficient training, or scalable serving, but not all three simultaneously. To address this limitation, we investigate sparse fine-tuning and observe a remarkable improvement in generalization ability. Utilizing this key insight, we propose a…

Cited by 3SourcePDFScholar
2024

SIRIUS : Contexual Sparisty with Correction for Efficient LLMs

NeurIPS 2024poster

With the blossom of large language models (LLM), inference efficiency becomes increasingly important. Various approximate methods are proposed to reduce the cost at inference time. Contextual Sparsity (CS) is appealing for its training-free nature and its ability to reach a higher compression ratio…

2024

Sequoia: Scalable and Robust Speculative Decoding

NeurIPS 2024spotlight

As the usage of large language models (LLMs) grows, it becomes increasingly important to serve them quickly and efficiently. While speculative decoding has recently emerged as a promising direction for accelerating LLM serving, existing methods are limited in their ability to scale to larger specula…

Cited by 3SourcePDFScholar
2024

Soft Prompt Recovers Compressed LLMs, Transferably

ICML 2024poster

Model compression is one of the most popular approaches to improve the accessibility of Large Language Models (LLMs) by reducing their memory footprint. However, the gaining of such efficiency benefits often simultaneously demands extensive engineering efforts and intricate designs to mitigate the p…

2024

SpecExec: Massively Parallel Speculative Decoding For Interactive LLM Inference on Consumer Devices

NeurIPS 2024poster

As large language models gain widespread adoption, running them efficiently becomes a crucial task. Recent works on LLM inference use speculative decoding to achieve extreme speedups. However, most of these works implicitly design their algorithms for high-end datacenter hardware. In this work, we a…

2023

CocktailSGD: Fine-tuning Foundation Models over 500Mbps Networks

ICML 2023poster

Distributed training of foundation models, especially large language models (LLMs), is communication-intensive and so has heavily relied on centralized data centers with fast interconnects. Can we train on slow networks and unlock the potential of decentralized infrastructure for foundation models?…

Cited by 49SourcePDFScholar
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

FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU

ICML 2023oral

The high computational and memory requirements of large language model (LLM) inference make it feasible only with multiple high-end accelerators. Motivated by the emerging demand for latency-insensitive tasks with batched processing, this paper initiates the study of high-throughput LLM inference us…

2023

H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models

NeurIPS 2023poster

Large Language Models (LLMs), despite their recent impressive accomplishments, are notably cost-prohibitive to deploy, particularly for applications involving long-content generation, such as dialogue systems and story writing. Often, a large amount of transient state information, referred to as the…

2023

Laughing Hyena Distillery: Extracting Compact Recurrences From Convolutions

NeurIPS 2023poster

Recent advances in attention-free sequence models rely on convolutions as alternatives to the attention operator at the core of Transformers. In particular, long convolution sequence models have achieved state-of-the-art performance in many domains, but incur a significant cost during auto-regressiv…

Cited by 25SourcePDFScholar
2023

Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

NeurIPS 2023poster

Transformer architecture has shown impressive performance in multiple research domains and has become the backbone of many neural network models. However, there is limited understanding on how it works. In particular, with a simple predictive loss, how the representation emerges from the gradient \…

Cited by 90SourcePDFScholar
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

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…

2021

A Tale of Two Efficient and Informative Negative Sampling Distributions

ICML 2021oral

Softmax classifiers with a very large number of classes naturally occur in many applications such as natural language processing and information retrieval. The calculation of full softmax is costly from the computational and energy perspective. There have been various sampling approaches to overcome…

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

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…

2018

LSH-SAMPLING BREAKS THE COMPUTATIONAL CHICKEN-AND-EGG LOOP IN ADAPTIVE STOCHASTIC GRADIENT ESTIMATION

ICLR 2018workshop

Stochastic Gradient Descent or SGD is the most popular optimization algorithm for large-scale problems. SGD estimates the gradient by uniform sampling with sample size one. There have been several other works that suggest faster epoch wise convergence by using weighted non-uniform sampling for bette…

Cited by 13SourceScholar