← Search

Jianfei Chen

46 accepted papers

2026

Achieving low-bit Muon through subspace preservation and grid quantization

ICLR 2026poster

Training Large Language Models (LLMs) faces severe memory constraints due to the increasing size of model parameters and optimizer states. The Muon optimizer, which is based on matrix orthogonalization, has recently demonstrated significant potential and offers considerable memory advantages over Ad…

Cited by 0SourcecodeScholar
2026

Deterministic Differentiable Structured Pruning for Large Language Models

ICML 2026poster

Structured pruning reduces LLM inference cost by removing low-importance architectural components. This can be viewed as learning a multiplicative gate for each component under an $\ell_0$ sparsity constraint. Due to the discreteness of the $\ell_0$ norm, prior work typically adopts stochastic hard-…

Cited by 0SourceScholar
2026

HexGen-3: A Fully Disaggregated LLM Serving Framework with Fine-Grained Heterogeneous Resource Autoscaling

ICML 2026poster

The operational cost of serving large language models remains prohibitively high, largely due to extreme workload heterogeneity in production traffic. We observe that combining disaggregated inference with resource autoscaling enables fine-grained resource adjustment, allowing inference phases and o…

Cited by 0SourceScholar
2026

Large Scale Diffusion Distillation via Score-Regularized Continuous-Time Consistency

ICLR 2026poster

Although continuous-time consistency models (e.g., sCM, MeanFlow) are theoretically principled and empirically powerful for fast academic-scale diffusion, its applicability to large-scale text-to-image and video tasks remains unclear due to infrastructure challenges in Jacobian-vector product (JVP)…

Cited by 0SourcecodeScholar
2026

SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse–Linear Attention

ICLR 2026poster

In Diffusion Transformer (DiT) models, particularly for video generation, attention latency is a major bottleneck due to the long sequence length and the quadratic complexity. Interestingly, we find that attention weights can be decoupled into two matrices: a small fraction of large weights with hig…

Cited by 44SourcecodeScholar
2026

TetraJet-v2: Accurate NVFP4 Training for Large Language Models with Oscillation Suppression and Outlier Control

ICML 2026spotlight

Large Language Models (LLMs) training is prohibitively expensive, driving interest in low-precision fully-quantized training (FQT). While novel 4-bit formats like NVFP4 offer substantial efficiency gains, achieving near-lossless training at such low precision remains challenging. We introduce **Tetr…

Cited by 0SourceScholar
2025

COAT: Compressing Optimizer states and Activations for Memory-Efficient FP8 Training

ICLR 2025poster

FP8 training has emerged as a promising method for improving training efficiency. Existing frameworks accelerate training by applying FP8 computation to linear layers while leaving optimizer states and activations in higher precision, which fails to fully optimize memory usage. This paper introduces…

Cited by 4SourcePDFScholar
2025

Elucidating the Preconditioning in Consistency Distillation

ICLR 2025poster

Consistency distillation is a prevalent way for accelerating diffusion models adopted in consistency (trajectory) models, in which a student model is trained to traverse backward on the probability flow (PF) ordinary differential equation (ODE) trajectory determined by the teacher model. Preconditio…

Cited by 2SourcePDFScholar
2025

FrameBridge: Improving Image-to-Video Generation with Bridge Models

ICML 2025poster

Diffusion models have achieved remarkable progress on image-to-video (I2V) generation, while their noise-to-data generation process is inherently mismatched with this task, which may lead to suboptimal synthesis quality. In this work, we present FrameBridge. By modeling the frame-to-frames generatio…

2025

On the Optimization and Generalization of Two-layer Transformers with Sign Gradient Descent

ICLR 2025spotlight

The Adam optimizer is widely used for transformer optimization in practice, which makes understanding the underlying optimization mechanisms an important problem. However, due to the Adam's complexity, theoretical analysis of how it optimizes transformers remains a challenging task. Fortunately, Si…

Cited by 3SourcePDFScholar
2025

Oscillation-Reduced MXFP4 Training for Vision Transformers

ICML 2025poster

Pre-training Transformers in FP4 precision is becoming a promising approach to gain substantial speedup, but it comes with a considerable loss of accuracy. Microscaling (MX) data format provides a fine-grained per-group quantization method to improve the representation ability of the FP4 format and…

2025

Pruning Large Language Models with Semi-Structural Adaptive Sparse Training

AAAI 2025technical

The remarkable success of Large Language Models (LLMs) relies heavily on their substantial scale, which poses significant challenges during model deployment in terms of latency and memory consumption. Recently, numerous studies have attempted to compress LLMs using one-shot pruning methods. However,…

2025

SageAttention2: Efficient Attention with Thorough Outlier Smoothing and Per-thread INT4 Quantization

ICML 2025poster

Although quantization for linear layers has been widely used, its application to accelerate the attention process remains limited. To further enhance the efficiency of attention computation compared to SageAttention while maintaining precision, we propose SageAttention2, which utilizes significantly…

2025

SageAttention3: Microscaling FP4 Attention for Inference and An Exploration of 8-Bit Training

NeurIPS 2025spotlight

The efficiency of attention is important due to its quadratic time complexity. We enhance the efficiency of attention through two key contributions: First, we leverage the new $\texttt{FP4}$ Tensor Cores in Blackwell GPUs to accelerate attention computation. Our implementation achieves $\textbf{1038…

Cited by 0SourcecodeScholar
2025

SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration

ICLR 2025poster

The transformer architecture predominates across various models. As the heart of the transformer, attention has a computational complexity of $O(N^2)$, compared to $O(N)$ for linear transformations. When handling large sequence lengths, attention becomes the primary time-consuming component. Althoug…

2025

SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference

ICML 2025poster

An efficient attention implementation is essential for large models due to its quadratic time complexity. Fortunately, attention commonly exhibits sparsity, i.e., many values in the attention map are near zero, allowing for the omission of corresponding computations. Many studies have utilized the s…

2025

Sparse Video-Gen: Accelerating Video Diffusion Transformers with Spatial-Temporal Sparsity

ICML 2025poster

Diffusion Transformers (DiTs) dominate video generation but their high computational cost severely limits real-world applicability, usually requiring tens of minutes to generate a few seconds of video even on high-performance GPUs. This inefficiency primarily arises from the quadratic computational…

Cited by 11SourcePDFScholar
2025

Sparse VideoGen2: Accelerate Video Generation with Sparse Attention via Semantic-Aware Permutation

NeurIPS 2025spotlight

Diffusion Transformers (DiTs) are essential for video generation but suffer from significant latency due to the quadratic complexity of attention. By computing only critical tokens, sparse attention reduces computational costs and offers a promising acceleration approach. However, we identify that…

Cited by 0SourcecodeScholar
2025

Task-Specific Zero-shot Quantization-Aware Training for Object Detection

ICCV 2025poster

Quantization is a key technique to reduce network size and computational complexity by representing the network parameters with a lower precision. Traditional quantization methods rely on access to original training data, which is often restricted due to privacy concerns or security challenges. Zero…

Cited by 0SourcePDFScholar
2024

Accelerating Transformer Pre-training with 2:4 Sparsity

ICML 2024poster

Training large transformers is slow, but recent innovations on GPU architecture give us an advantage. NVIDIA Ampere GPUs can execute a fine-grained 2:4 sparse matrix multiplication twice as fast as its dense equivalent. In the light of this property, we comprehensively investigate the feasibility of…

2024

C-GAIL: Stabilizing Generative Adversarial Imitation Learning with Control Theory

NeurIPS 2024poster

Generative Adversarial Imitation Learning (GAIL) provides a promising approach to training a generative policy to imitate a demonstrator. It uses on-policy Reinforcement Learning (RL) to optimize a reward signal derived from an adversarial discriminator. However, optimizing GAIL is difficult in prac…

Cited by 2SourcePDFScholar
2024

Jetfire: Efficient and Accurate Transformer Pretraining with INT8 Data Flow and Per-Block Quantization

ICML 2024spotlight

Pretraining transformers are generally time-consuming. Fully quantized training (FQT) is a promising approach to speed up pretraining. However, most FQT methods adopt a quantize-compute-dequantize procedure, which often leads to suboptimal speedup and significant performance degradation when used in…

2023

Contrastive Energy Prediction for Exact Energy-Guided Diffusion Sampling in Offline Reinforcement Learning

ICML 2023poster

Guided sampling is a vital approach for applying diffusion models in real-world tasks that embeds human-defined guidance during the sampling procedure. This paper considers a general setting where the guidance is defined by an (unnormalized) energy function. The main challenge for this setting is th…

2023

DPM-Solver-v3: Improved Diffusion ODE Solver with Empirical Model Statistics

NeurIPS 2023poster

Diffusion probabilistic models (DPMs) have exhibited excellent performance for high-fidelity image generation while suffering from inefficient sampling. Recent works accelerate the sampling procedure by proposing fast ODE solvers that leverage the specific ODE form of DPMs. However, they highly rely…

2023

Improved Techniques for Maximum Likelihood Estimation for Diffusion ODEs

ICML 2023poster

Diffusion models have exhibited excellent performance in various domains. The probability flow ordinary differential equation (ODE) of diffusion models (i.e., diffusion ODEs) is a particular case of continuous normalizing flows (CNFs), which enables deterministic inference and exact likelihood evalu…

2023

Preserving Pre-trained Features Helps Calibrate Fine-tuned Language Models

ICLR 2023poster

Large pre-trained language models (PLMs) have demonstrated strong performance on natural language understanding (NLU) tasks through fine-tuning. However, fine-tuned models still suffer from overconfident predictions, especially in out-of-domain settings. In this paper, we tackle the problem of calib…

2022

DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps

NeurIPS 2022accept

Diffusion probabilistic models (DPMs) are emerging powerful generative models. Despite their high-quality generation performance, DPMs still suffer from their slow sampling as they generally need hundreds or thousands of sequential function evaluations (steps) of large neural networks to draw a samp…

2022

Fast Lossless Neural Compression with Integer-Only Discrete Flows

ICML 2022spotlight

By applying entropy codecs with learned data distributions, neural compressors have significantly outperformed traditional codecs in terms of compression ratio. However, the high inference latency of neural networks hinders the deployment of neural compressors in practical applications. In this work…

2022

GACT: Activation Compressed Training for Generic Network Architectures

ICML 2022spotlight

Training large neural network (NN) models requires extensive memory resources, and Activation Compression Training (ACT) is a promising approach to reduce training memory footprint. This paper presents GACT, an ACT framework to support a broad range of machine learning tasks for generic NN architect…

2022

Maximum Likelihood Training for Score-based Diffusion ODEs by High Order Denoising Score Matching

ICML 2022spotlight

Score-based generative models have excellent performance in terms of generation quality and likelihood. They model the data distribution by matching a parameterized score network with first-order data score functions. The score network can be used to define an ODE (“score-based diffusion ODE”) for e…

2021

ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training

ICML 2021oral

The increasing size of neural network models has been critical for improvements in their accuracy, but device memory is not growing at the same rate. This creates fundamental challenges for training neural networks within limited memory environments. In this work, we propose ActNN, a memory-efficien…

2020

A Statistical Framework for Low-bitwidth Training of Deep Neural Networks

NeurIPS 2020poster

Fully quantized training (FQT), which uses low-bitwidth hardware by quantizing the activations, weights, and gradients of a neural network model, is a promising approach to accelerate the training of deep neural networks. One major challenge with FQT is the lack of theoretical understanding, in part…

2020

VFlow: More Expressive Generative Flows with Variational Data Augmentation

ICML 2020poster

Generative flows are promising tractable models for density modeling that define probabilistic distributions with invertible transformations. However, tractability imposes architectural constraints on generative flows. In this work, we study a previously overlooked constraint that all the intermedia…

2018

Stochastic Expectation Maximization with Variance Reduction

NeurIPS 2018poster

Expectation-Maximization (EM) is a popular tool for learning latent variable models, but the vanilla batch EM does not scale to large data sets because the whole data set is needed at every E-step. Stochastic Expectation Maximization (sEM) reduces the cost of E-step by stochastic approximation. Howe…

2018

Stochastic Training of Graph Convolutional Networks with Variance Reduction

ICML 2018oral

Graph convolutional networks (GCNs) are powerful deep neural networks for graph-structured data. However, GCN computes the representation of a node recursively from its neighbors, making the receptive field size grow exponentially with the number of layers. Previous attempts on reducing the receptiv…

2017

Population Matching Discrepancy and Applications in Deep Learning

NeurIPS 2017poster

A differentiable estimation of the distance between two distributions based on samples is important for many deep learning tasks. One such estimation is maximum mean discrepancy (MMD). However, MMD suffers from its sensitive kernel bandwidth hyper-parameter, weak gradients, and large mini-batch siz…

Cited by 7SourcePDFScholar