← Search

Zhewei Yao

32 accepted papers

2026

Agent World Model: Infinity Synthetic Environments for Agentic Reinforcement Learning

ICML 2026poster

Recent advances in large language model (LLM) have empowered autonomous agents to perform complex tasks that require multi-turn interactions with external tools and environments. However, scaling such agent training is limited by the lack of diverse and reliable environments. In this paper, we propo…

Cited by 0SourceScholar
2026

DARE-bench: Evaluating Modeling and Instruction Fidelity of LLMs in Data Science

ICLR 2026poster

The fast-growing demands in using Large Language Models (LLMs) to tackle complex multi-step data science tasks create a emergent need for accurate benchmarking. There are two major gaps in existing benchmarks: (i) the lack of standardized, process-aware evaluation that captures instruction adherence…

Cited by 0SourcecodeScholar
2025

CORD: Balancing COnsistency and Rank Distillation for Robust Retrieval-Augmented Generation

NAACL 2025short

With the adoption of retrieval-augmented generation (RAG), large language models (LLMs) are expected to ground their generation to the retrieved contexts. Yet, this is hindered by position bias of LLMs, failing to evenly attend to all contexts. Previous work has addressed this by synthesizing contex…

2025

Inference Scaling for Bridging Retrieval and Augmented Generation

NAACL 2025findings

Retrieval-augmented generation (RAG) has emerged as a popular approach to steering the output of a large language model (LLM) by incorporating retrieved contexts as inputs. However, existing work observed the generator bias, such that improving the retrieval results may negatively affect the outcome…

2025

STUN: Structured-Then-Unstructured Pruning for Scalable MoE Pruning

ACL 2025long

Mixture-of-experts (MoEs) have been adopted for reducing inference costs by sparsely activating experts in large language models (LLMs). Despite these reductions, the massive number of parameters in MoEs still makes them expensive to serve. Conventionally, unstructured or structured pruning has been…

2025

SwiftKV: Fast Prefill-Optimized Inference with Knowledge-Preserving Model Transformation

EMNLP 2025

LLM inference for enterprise applications, such as summarization, RAG, and code-generation, typically observe much longer prompt than generations, leading to high prefill cost and response latency. We present SwiftKV, a novel model transformation and distillation procedure targeted at reducing the p

2024

DeepSpeed Data Efficiency: Improving Deep Learning Model Quality and Training Efficiency via Efficient Data Sampling and Routing

AAAI 2024technical

Recent advances on deep learning models come at the price of formidable training cost. The increasing model size is one of the root causes, but another less-emphasized fact is that data scale is actually increasing at a similar speed as model scale, and the training cost is proportional to both of t…

2024

Exploring Post-training Quantization in LLMs from Comprehensive Study to Low Rank Compensation

AAAI 2024technical

Post-training quantization (PTQ) has emerged as a promising technique for mitigating memory consumption and computational costs in large language models (LLMs). However, a systematic examination of various quantization schemes, model families, and quantization bit precision has been absent from the…

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

ZeRO++: Extremely Efficient Collective Communication for Large Model Training

ICLR 2024poster

Zero Redundancy Optimizer (ZeRO) has been used to train a wide range of large language models on massive GPU clusters due to its ease of use, efficiency, and good scalability. However, when training on low-bandwidth clusters, and/or when small batch size per GPU is used, ZeRO’s effective throughput…

Cited by 9SourcePDFScholar
2023

DySR: Adaptive Super-Resolution via Algorithm and System Co-design

ICLR 2023poster

Super resolution (SR) is a promising approach for improving the quality of low resolution steaming services on mobile devices. On mobile devices, the available computing and memory resources change dynamically depending on other running applications. Due to the high computation and memory demands of…

Cited by 1SourcePDFScholar
2023

Scaling Vision-Language Models with Sparse Mixture of Experts

EMNLP 2023long findings

The field of natural language processing (NLP) has made significant strides in recent years, particularly in the development of large-scale vision-language models (VLMs). These models aim to bridge the gap between text and visual information, enabling a more comprehensive understanding of multimedia…

Cited by 0SourceScholar
2023

Understanding Int4 Quantization for Language Models: Latency Speedup, Composability, and Failure Cases

ICML 2023poster

Improving the deployment efficiency of transformer-based language models has been challenging given their high computation and memory cost. While INT8 quantization has recently been shown to be effective in reducing both the memory cost and latency while preserving model accuracy, it remains unclear…

Cited by 51SourcePDFScholar
2022

DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale

ICML 2022spotlight

As the training of giant dense models hits the boundary on the availability and capability of the hardware resources today, Mixture-of-Experts (MoE) models have become one of the most promising model architectures due to their significant training cost reduction compared to quality-equivalent dense…

2022

How Much Can CLIP Benefit Vision-and-Language Tasks?

ICLR 2022poster

Most existing Vision-and-Language (V&L) models rely on pre-trained visual encoders, using a relatively small set of manually-annotated data (as compared to web-crawled data), to perceive the visual world. However, it has been observed that large-scale pretraining usually can result in better general…

2022

Integer-Only Zero-Shot Quantization for Efficient Speech Recognition

ICASSP 2022accepted

End-to-end neural network models achieve improved performance on various automatic speech recognition (ASR) tasks. However, these models perform poorly on edge hardware due to large memory and computation requirements. While quantizing model weights and/or activations to low-precision can be a promi…

Cited by 0SourceScholar
2022

XTC: Extreme Compression for Pre-trained Transformers Made Simple and Efficient

NeurIPS 2022accept

Extreme compression, particularly ultra-low bit precision (binary/ternary) quantization, has been proposed to fit large NLP models on resource-constraint devices. However, to preserve the accuracy for such aggressive compression schemes, cutting-edge methods usually introduce complicated compressio…

Cited by 29SourcePDFScholar
2022

ZeroQuant: Efficient and Affordable Post-Training Quantization for Large-Scale Transformers

NeurIPS 2022accept

How to efficiently serve ever-larger trained natural language models in practice has become exceptionally challenging even for powerful cloud servers due to their prohibitive memory/computation requirements. In this work, we present an efficient and affordable post-training quantization approach to…

2021

ADAHESSIAN: An Adaptive Second Order Optimizer for Machine Learning

AAAI 2021technical

Incorporating second-order curvature information into machine learning optimization algorithms can be subtle, and doing so naïvely can lead to high per-iteration costs associated with forming the Hessian and performing the associated linear system solve. To address this, we introduce ADAHESSIAN, a n…

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…

2021

HAWQ-V3: Dyadic Neural Network Quantization

ICML 2021spotlight

Current low-precision quantization algorithms often have the hidden cost of conversion back and forth from floating point to quantized integer values. This hidden cost limits the latency improvement realized by quantizing Neural Networks. To address this, we present HAWQ-V3, a novel mixed-precision…

2021

I-BERT: Integer-only BERT Quantization

ICML 2021oral

Transformer based models, like BERT and RoBERTa, have achieved state-of-the-art results in many Natural Language Processing tasks. However, their memory footprint, inference latency, and power consumption are prohibitive efficient inference at the edge, and even at the data center. While quantizatio…

2021

What’s Hidden in a One-layer Randomly Weighted Transformer?

EMNLP 2021main

We demonstrate that, hidden within one-layer randomly weighted neural networks, there exist subnetworks that can achieve impressive performance, without ever modifying the weight initializations, on machine translation tasks. To find subnetworks for one-layer randomly weighted neural networks, we ap…

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

HAWQ-V2: Hessian Aware trace-Weighted Quantization of Neural Networks

NeurIPS 2020poster

Quantization is an effective method for reducing memory footprint and inference time of Neural Networks. However, ultra low precision quantization could lead to significant degradation in model accuracy. A promising method to address this is to perform mixed-precision quantization, where more sensit…

Cited by 338SourcePDFScholar
2020

PowerNorm: Rethinking Batch Normalization in Transformers

ICML 2020poster

The standard normalization method for neural network (NN) models used in Natural Language Processing (NLP) is layer normalization (LN).This is different than batch normalization (BN), which is widely-adopted in Computer Vision. The preferred use of LN in NLP is principally due to the empirical obser…

2020

ZeroQ: A Novel Zero Shot Quantization Framework

CVPR 2020poster

Quantization is a promising approach for reducing the inference time and memory footprint of neural networks. However, most existing quantization methods require access to the original training dataset for retraining during quantization. This is often not possible for applications with sensitive or…

Cited by 513PDFcodeScholar
2019

ANODEV2: A Coupled Neural ODE Framework

NeurIPS 2019poster

It has been observed that residual networks can be viewed as the explicit Euler discretization of an Ordinary Differential Equation (ODE). This observation motivated the introduction of so-called Neural ODEs, in which other discretization schemes and/or adaptive time stepping techniques can be used…

2019

HAWQ: Hessian AWare Quantization of Neural Networks With Mixed-Precision

ICCV 2019poster

Model size and inference speed/power have become a major challenge in the deployment of neural networks for many applications. A promising approach to address these problems is quantization. However, uniformly quantizing a model to ultra-low precision leads to significant accuracy degradation. A nov…

Cited by 637PDFcodeScholar
2019

Trust Region Based Adversarial Attack on Neural Networks

CVPR 2019poster

Deep Neural Networks are quite vulnerable to adversarial perturbations. Current state-of-the-art adversarial attack methods typically require very time consuming hyper-parameter tuning, or require many iterations to solve an optimization based adversarial attack. To address this problem, we present…

Cited by 78PDFcodeScholar
2018

Hessian-based Analysis of Large Batch Training and Robustness to Adversaries

NeurIPS 2018poster

Large batch size training of Neural Networks has been shown to incur accuracy loss when trained with the current methods. The exact underlying reasons for this are still not completely understood. Here, we study large batch size training through the lens of the Hessian operator and robust optimiza…