← Search

Yu-Xiang Wang

101 accepted papers

2026

Generalization Below the Edge of Stability: The Role of Data Geometry

ICLR 2026poster

Understanding generalization in overparameterized neural networks hinges on the interplay between the data geometry, neural architecture, and training dynamics. In this paper, we theoretically explore how data geometry controls this implicit bias. This paper presents theoretical results for overpara…

Cited by 0SourceScholar
2026

Not-a-Bandit: Provably No-Regret Drafter Selection in Speculative Decoding for LLMs

ICLR 2026poster

Speculative decoding is widely used in accelerating large language model (LLM) inference. In this work, we focus on the online draft model selection problem in speculative decoding. We design an algorithm that provably competes with the best draft model in hindsight for each query in terms of eithe…

Cited by 0SourceScholar
2025

A Technical Report on “Erasing the Invisible”: The 2024 NeurIPS Competition on Stress Testing Image Watermarks

NeurIPS 2025poster

AI-generated images have become pervasive, raising critical concerns around content authenticity, intellectual property, and the spread of misinformation. Invisible watermarks offer a promising solution for identifying AI-generated images, preserving content provenance without degrading visual quali…

Cited by 0SourceScholar
2025

Adapting to Linear Separable Subsets with Large-Margin in Differentially Private Learning

ICML 2025poster

This paper studies the problem of differentially private empirical risk minimization (DP-ERM) for binary linear classification. We obtain an efficient $(\varepsilon,\delta)$-DP algorithm with an empirical zero-one risk bound of $\tilde{O}\left(\frac{1}{\gamma^2\varepsilon n} + \frac{|S_{\mathrm{o…

Cited by 0SourcePDFScholar
2025

Adapting to Online Distribution Shifts in Deep Learning: A Black-Box Approach

AISTATS 2025poster

We study the well-motivated problem of online distribution shift in which the data arrive in batches and the distribution of each batch can change arbitrarily over time. Since the shifts can be large or small, abrupt or gradual, the length of the relevant historical data to learn from may vary over…

Cited by 0SourceScholar
2025

Adaptive Estimation and Learning under Temporal Distribution Shift

ICML 2025poster

In this paper, we study the problem of estimation and learning under temporal distribution shift. Consider an observation sequence of length $n$, which is a noisy realization of a time-varying ground-truth sequence. Our focus is to develop methods to estimate the groundtruth at the final time-step w…

Cited by 0SourcePDFScholar
2025

Efficiently Identifying Watermarked Segments in Mixed-Source Texts

ACL 2025long

Text watermarks in large language models (LLMs) are increasingly used to detect synthetic text, mitigating misuse cases like fake news and academic dishonesty. While existing watermarking detection techniques primarily focus on classifying entire documents as watermarked or not, they often neglect t…

2025

PROXSPARSE: REGULARIZED LEARNING OF SEMI-STRUCTURED SPARSITY MASKS FOR PRETRAINED LLMS

ICML 2025poster

Large Language Models (LLMs) have demonstrated exceptional performance in natural language processing tasks, yet their massive size makes serving them inefficient and costly. Semi-structured pruning has emerged as an effective method for model acceleration, but existing approaches are suboptimal bec…

Cited by 0SourcePDFScholar
2025

Permute-and-Flip: An optimally stable and watermarkable decoder for LLMs

ICLR 2025poster

In this paper, we propose a new decoding method called Permute-and-Flip (PF) decoder. It enjoys stability properties similar to the standard sampling decoder, but is provably up to 2x better in its quality-stability tradeoff than sampling and never worse than any other decoder. We also design a cryp…

2025

Purifying Approximate Differential Privacy with Randomized Post-processing

NeurIPS 2025spotlight

We propose a framework to convert $(\varepsilon, \delta)$-approximate Differential Privacy (DP) mechanisms into $(\varepsilon', 0)$-pure DP mechanisms under certain conditions, a process we call ``purification.'' This algorithmic technique leverages randomized post-processing with calibrated noise t…

Cited by 0SourceScholar
2025

Stable Minima of ReLU Neural Networks Suffer from the Curse of Dimensionality: The Neural Shattering Phenomenon

NeurIPS 2025spotlight

We study the implicit bias of flatness / low (loss) curvature and its effects on generalization in two-layer overparameterized ReLU networks with multivariate inputs---a problem well motivated by the minima stability and edge-of-stability phenomena in gradient-descent training. Existing work either…

Cited by 0SourceScholar
2025

Weak-to-Strong Jailbreaking on Large Language Models

ICML 2025poster

Large language models (LLMs) are vulnerable to jailbreak attacks -- resulting in harmful, unethical, or biased text generations. However, existing jailbreaking methods are computationally costly. In this paper, we propose the **weak-to-strong** jailbreaking attack, an efficient inference time attack…

2024

CPR: Retrieval Augmented Generation for Copyright Protection

CVPR 2024poster

Retrieval Augmented Generation (RAG) is emerging as a flexible and robust technique to adapt models to private users data without training to handle credit attribution and to allow efficient machine unlearning at scale. However RAG techniques for image generation may lead to parts of the retrieved s…

Cited by 79SourcePDFScholar
2024

Differentially Private Bias-Term Fine-tuning of Foundation Models

ICML 2024poster

We study the problem of differentially private (DP) fine-tuning of large pre-trained models — a recent privacy-preserving approach suitable for solving downstream tasks with sensitive data. Existing work has demonstrated that high accuracy is possible under strong privacy constraint, yet requires si…

2024

Improving Sample Efficiency of Model-Free Algorithms for Zero-Sum Markov Games

ICML 2024poster

The problem of two-player zero-sum Markov games has recently attracted increasing interests in theoretical studies of multi-agent reinforcement learning (RL). In particular, for finite-horizon episodic Markov decision processes (MDPs), it has been shown that model-based algorithms can find an $\epsi…

Cited by 1SourcePDFScholar
2024

Invisible Image Watermarks Are Provably Removable Using Generative AI

NeurIPS 2024poster

Invisible watermarks safeguard images' copyrights by embedding hidden messages only detectable by owners. They also prevent people from misusing images, especially those generated by AI models. We propose a family of regeneration attacks to remove these invisible watermarks. The proposed attack met…

2024

NetworkGym: Reinforcement Learning Environments for Multi-Access Traffic Management in Network Simulation

NeurIPS 2024poster

Mobile devices such as smartphones, laptops, and tablets can often connect to multiple access networks (e.g., Wi-Fi, LTE, and 5G) simultaneously. Recent advancements facilitate seamless integration of these connections below the transport layer, enhancing the experience for apps that lack inherent m…

2024

Neural Collapse meets Differential Privacy: Curious behaviors of NoisyGD with Near-Perfect Representation Learning

ICML 2024oral

A recent study by De et al. (2022) shows that large-scale representation learning through pre-training on a public dataset significantly enhances differentially private (DP) learning in downstream tasks. To explain this, we consider a layer-peeled model in representation learning, resulting in Neura…

Cited by 0SourcePDFScholar
2024

Nonparametric Classification on Low Dimensional Manifolds using Overparameterized Convolutional Residual Networks

NeurIPS 2024poster

Convolutional residual neural networks (ConvResNets), though overparametersized, can achieve remarkable prediction performance in practice, which cannot be well explained by conventional wisdom. To bridge this gap, we study the performance of ConvResNeXts trained with weight decay, which cover ConvR…

Cited by 2SourcePDFScholar
2024

Online Feature Updates Improve Online (Generalized) Label Shift Adaptation

NeurIPS 2024poster

This paper addresses the prevalent issue of label shift in an online setting with missing labels, where data distributions change over time and obtaining timely labels is challenging. While existing methods primarily focus on adjusting or updating the final layer of a pre-trained classifier, we expl…

Cited by 2SourcePDFScholar
2024

Provable Robust Watermarking for AI-Generated Text

ICLR 2024poster

We study the problem of watermarking large language models (LLMs) generated text — one of the most promising approaches for addressing the safety challenges of LLM usage. In this paper, we propose a rigorous theoretical framework to quantify the effectiveness and robustness of LLM watermarks. We pro…

Cited by 166SourcePDFScholar
2024

Stable Minima Cannot Overfit in Univariate ReLU Networks: Generalization by Large Step Sizes

NeurIPS 2024spotlight

We study the generalization of two-layer ReLU neural networks in a univariate nonparametric regression problem with noisy labels. This is a problem where kernels (\emph{e.g.} NTK) are provably sub-optimal and benign overfitting does not happen, thus disqualifying existing theory for interpolating (0…

Cited by 4SourcePDFScholar
2024

Tractable MCMC for Private Learning with Pure and Gaussian Differential Privacy

ICLR 2024poster

Posterior sampling, i.e., exponential mechanism to sample from the posterior distribution, provides $\varepsilon$-pure differential privacy (DP) guarantees and does not suffer from potentially unbounded privacy breach introduced by $(\varepsilon,\delta)$-approximate DP. In practice, however, one nee…

Cited by 6SourcePDFScholar
2023

A Privacy-Friendly Approach to Data Valuation

NeurIPS 2023spotlight

Data valuation, a growing field that aims at quantifying the usefulness of individual data sources for training machine learning (ML) models, faces notable yet often overlooked privacy challenges. This paper studies these challenges with a focus on KNN-Shapley, one of the most practical data valuati…

Cited by 24SourcePDFScholar
2023

Automatic Clipping: Differentially Private Deep Learning Made Easier and Stronger

NeurIPS 2023poster

Per-example gradient clipping is a key algorithmic step that enables practical differential private (DP) training for deep learning models. The choice of clipping threshold $R$, however, is vital for achieving high accuracy under DP. We propose an easy-to-use replacement, called automatic clipping,…

2023

Deep Learning meets Nonparametric Regression: Are Weight-Decayed DNNs Locally Adaptive?

ICLR 2023poster

We study the theory of neural network (NN) from the lens of classical nonparametric regression problems with a focus on NN’s ability to adaptively estimate functions with heterogeneous smoothness — a property of functions in Besov or Bounded Variation (BV) classes. Existing work on this problem requ…

Cited by 18SourcePDFScholar
2023

Differentially Private Optimization on Large Model at Small Cost

ICML 2023poster

Differentially private (DP) optimization is the standard paradigm to learn large neural networks that are accurate and privacy-preserving. The computational cost for DP deep learning, however, is notoriously heavy due to the per-sample gradient clipping. Existing DP implementations are 2$\sim$1000$…

2023

Generalized PTR: User-Friendly Recipes for Data-Adaptive Algorithms with Differential Privacy

AISTATS 2023poster

The “Propose-Test-Release” (PTR) framework [Dwork and Lei, 2009] is a classic recipe for designing differentially private (DP) algorithms that are data-adaptive, i.e. those that add less noise when the input dataset is “nice”. We extend PTR to a more general setting by privately testing data-depende…

Cited by 8SourcePDFScholar
2023

Improving the Privacy and Practicality of Objective Perturbation for Differentially Private Linear Learners

NeurIPS 2023poster

In the arena of privacy-preserving machine learning, differentially private stochastic gradient descent (DP-SGD) has outstripped the objective perturbation mechanism in popularity and interest. Though unrivaled in versatility, DP-SGD requires a non-trivial privacy overhead (for privately tuning the…

Cited by 10SourcePDFScholar
2023

Near-Optimal Deployment Efficiency in Reward-Free Reinforcement Learning with Linear Function Approximation

ICLR 2023poster

We study the problem of deployment efficient reinforcement learning (RL) with linear function approximation under the \emph{reward-free} exploration setting. This is a well-motivated problem because deploying new policies is costly in real-life RL applications. Under the linear MDP setting with feat…

Cited by 13SourcePDFScholar
2023

Non-stationary Reinforcement Learning under General Function Approximation

ICML 2023poster

General function approximation is a powerful tool to handle large state and action spaces in a broad range of reinforcement learning (RL) scenarios. However, theoretical understanding of non-stationary MDPs with general function approximation is still limited. In this paper, we make the first such a…

Cited by 8SourcePDFScholar
2023

Offline Reinforcement Learning with Closed-Form Policy Improvement Operators

ICML 2023poster

Behavior constrained policy optimization has been demonstrated to be a successful paradigm for tackling Offline Reinforcement Learning. By exploiting historical transitions, a policy is trained to maximize a learned value function while constrained by the behavior policy to avoid a significant distr…

2023

Offline Reinforcement Learning with Differentiable Function Approximation is Provably Efficient

ICLR 2023poster

Offline reinforcement learning, which aims at optimizing sequential decision-making strategies with historical data, has been extensively applied in real-life applications. State-Of-The-Art algorithms usually leverage powerful function approximators (e.g. neural networks) to alleviate the sample com…

Cited by 19SourcePDFScholar
2023

Online Label Shift: Optimal Dynamic Regret meets Practical Algorithms

NeurIPS 2023spotlight

This paper focuses on supervised and unsupervised online label shift, where the class marginals $Q(y)$ varies but the class-conditionals $Q(x|y)$ remain invariant. In the unsupervised setting, our goal is to adapt a learner, trained on some offline labeled data, to changing label distributions given…

2023

Posterior Sampling with Delayed Feedback for Reinforcement Learning with Linear Function Approximation

NeurIPS 2023poster

Recent studies in reinforcement learning (RL) have made significant progress by leveraging function approximation to alleviate the sample complexity hurdle for better performance. Despite the success, existing provably efficient algorithms typically rely on the accessibility of immediate feedback up…

Cited by 8SourcePDFScholar
2022

Adaptive Private-K-Selection with Adaptive K and Application to Multi-label PATE

AISTATS 2022poster

We provide an end-to-end Renyi DP based-framework for differentially private top-$k$ selection. Unlike previous approaches, which require a data-independent choice on $k$, we propose to privately release a data-dependent choice of $k$ such that the gap between $k$-th and the $(k+1)$st “quality” is l…

Cited by 16SourcePDFScholar
2022

Differentially Private Linear Sketches: Efficient Implementations and Applications

NeurIPS 2022accept

Linear sketches have been widely adopted to process fast data streams, and they can be used to accurately answer frequency estimation, approximate top K items, and summarize data distributions. When data are sensitive, it is desirable to provide privacy guarantees for linear sketches to preserve pri…

Cited by 31SourcePDFScholar
2022

Mixed Differential Privacy in Computer Vision

CVPR 2022oral

We introduce AdaMix, an adaptive differentially private algorithm for training deep neural network classifiers using both private and public image data. While pre-training language models on large public datasets has enabled strong differential privacy (DP) guarantees with minor loss of accuracy, a…

Cited by 64PDFcodeScholar
2022

Near-optimal Offline Reinforcement Learning with Linear Representation: Leveraging Variance Information with Pessimism

ICLR 2022poster

Offline reinforcement learning, which seeks to utilize offline/historical data to optimize sequential decision-making strategies, has gained surging prominence in recent studies. Due to the advantage that appropriate function approximators can help mitigate the sample complexity burden in modern rei…

Cited by 87SourcePDFScholar
2022

Offline stochastic shortest path: Learning, evaluation and towards optimality

UAI 2022poster

Goal-oriented Reinforcement Learning, where the agent needs to reach the goal state while simultaneously minimizing the cost, has received significant attention in real-world applications. Its theoretical formulation, stochastic shortest path (SSP), has been intensively researched in the online sett…

Cited by 7SourcePDFScholar
2022

Optimal Accounting of Differential Privacy via Characteristic Function

AISTATS 2022poster

Characterizing the privacy degradation over compositions, i.e., privacy accounting, is a fundamental topic in differential privacy (DP) with many applications to differentially private machine learning and federated learning. We propose a unification of recent advances (Renyi DP, privacy profiles, $…

2022

Sample-Efficient Reinforcement Learning with loglog(T) Switching Cost

ICML 2022spotlight

We study the problem of reinforcement learning (RL) with low (policy) switching cost {—} a problem well-motivated by real-life RL applications in which deployments of new policies are costly and the number of policy updates must be low. In this paper, we propose a new algorithm based on stage-wise e…

Cited by 37SourcePDFScholar
2022

SeqPATE: Differentially Private Text Generation via Knowledge Distillation

NeurIPS 2022accept

Protecting the privacy of user data is crucial for text generation models, which can leak sensitive information during generation. Differentially private (DP) learning methods provide guarantees against identifying the existence of a training sample from model outputs. PATE is a recent DP learning a…

Cited by 24SourcePDFScholar
2022

Towards Agnostic Feature-based Dynamic Pricing: Linear Policies vs Linear Valuation with Unknown Noise

AISTATS 2022poster

In feature-based dynamic pricing, a seller sets appropriate prices for a sequence of products (described by feature vectors) on the fly by learning from the binary outcomes of previous sales sessions ("Sold" if valuation $\geq$ price, and "Not Sold" otherwise). Existing works either assume noiseless…

Cited by 28SourcePDFScholar
2021

Near-Optimal Provable Uniform Convergence in Offline Policy Evaluation for Reinforcement Learning

AISTATS 2021poster

The problem of \emph{Offline Policy Evaluation} (OPE) in Reinforcement Learning (RL) is a critical step towards applying RL in real life applications. Existing work on OPE mostly focus on evaluating a \emph{fixed} target policy $\pi$, which does not provide useful bounds for offline policy learning…

Cited by 84SourcePDFScholar
2021

Optimal Uniform OPE and Model-based Offline Reinforcement Learning in Time-Homogeneous, Reward-Free and Task-Agnostic Settings

NeurIPS 2021poster

This work studies the statistical limits of uniform convergence for offline policy evaluation (OPE) problems with model-based methods (for episodic MDP) and provides a unified framework towards optimal learning for several well-motivated offline tasks. Uniform OPE $\sup_\Pi|Q^\pi-\hat{Q}^\pi|<\epsil…

Cited by 29SourcePDFScholar
2021

Revisiting Model-Agnostic Private Learning: Faster Rates and Active Learning

AISTATS 2021poster

The Private Aggregation of Teacher Ensembles (PATE) framework is one of the most promising recent approaches in differentially private learning. Existing theoretical analysis shows that PATE consistently learns any VC-classes in the realizable setting, but falls short in explaining its success in mo…

Cited by 17SourcePDFScholar
2020

An end-to-end Differentially Private Latent Dirichlet Allocation Using a Spectral Algorithm

ICML 2020poster

We provide an end-to-end differentially private spectral algorithm for learning LDA, based on matrix/tensor decompositions, and establish theoretical guarantees on utility/consistency of the estimated model parameters. We represent the spectral algorithm as a computational graph. Noise can be inject…

Cited by 12SourcePDFScholar
2020

Domain Adaptation with Conditional Distribution Matching and Generalized Label Shift

NeurIPS 2020poster

Adversarial learning has demonstrated good performance in the unsupervised domain adaptation setting, by learning domain-invariant representations. However, recent work has shown limitations of this approach when label distributions differ between the source and target domains. In this paper, we pro…

Cited by 214SourcePDFScholar
2019

A Higher-Order Kolmogorov-Smirnov Test

AISTATS 2019poster

We present an extension of the Kolmogorov-Smirnov (KS) two-sample test, which can be more sensitive to differences in the tails. Our test statistic is an integral probability metric (IPM) defined over a higher-order total variation ball, recovering the original KS test as its simplest case. We giv…

Cited by 18SourcePDFScholar
2019

Enhancing the Locality and Breaking the Memory Bottleneck of Transformer on Time Series Forecasting

NeurIPS 2019poster

Time series forecasting is an important problem across many domains, including predictions of solar plant energy output, electricity consumption, and traffic jam situation. In this paper, we propose to tackle such forecasting problem with Transformer. Although impressed by its performance in our pre…

Cited by 2073SourcePDFScholar
2019

Subsampled Renyi Differential Privacy and Analytical Moments Accountant

AISTATS 2019poster

We study the problem of subsampling in differential privacy (DP), a question that is the centerpiece behind many successful differentially private machine learning algorithms. Specifically, we provide a tight upper bound on the Renyi Differential Privacy (RDP) [Mironov 2017] parameters for algorith…

Cited by 463SourcePDFScholar
2019

Towards Optimal Off-Policy Evaluation for Reinforcement Learning with Marginalized Importance Sampling

NeurIPS 2019poster

Motivated by the many real-world applications of reinforcement learning (RL) that require safe-policy iterations, we consider the problem of off-policy evaluation (OPE) --- the problem of evaluating a new policy using the historical data obtained by different behavior policies --- under the model o…

Cited by 206SourcePDFScholar
2018

Detecting and Correcting for Label Shift with Black Box Predictors

ICML 2018oral

Faced with distribution shift between training and test set, we wish to detect and quantify the shift, and to correct our classifiers without test set labels. Motivated by medical diagnosis, where diseases (targets), cause symptoms (observations), we focus on label shift, where the label marginal p(…

2018

Improving the Gaussian Mechanism for Differential Privacy: Analytical Calibration and Optimal Denoising

ICML 2018oral

The Gaussian mechanism is an essential building block used in multitude of differentially private data analysis algorithms. In this paper we revisit the Gaussian mechanism and show that the original analysis has several important limitations. Our analysis reveals that the variance formula for the or…

2018

signSGD: Compressed Optimisation for Non-Convex Problems

ICML 2018oral

Training large neural networks requires distributing learning across multiple workers, where the cost of communicating gradients can be a significant bottleneck. signSGD alleviates this problem by transmitting just the sign of each minibatch stochastic gradient. We prove that it can get the best of…

2017

Higher-Order Total Variation Classes on Grids: Minimax Theory and Trend Filtering Methods

NeurIPS 2017poster

We consider the problem of estimating the values of a function over $n$ nodes of a $d$-dimensional grid graph (having equal side lengths $n^{1/d}$) from noisy observations. The function is assumed to be smooth, but is allowed to exhibit different amounts of smoothness at different regions in the gri…

Cited by 39SourcePDFScholar
2017

Optimal and Adaptive Off-policy Evaluation in Contextual Bandits

ICML 2017poster

We study the off-policy evaluation problem—estimating the value of a target policy using data collected by another policy—under the contextual bandit model. We consider the general (agnostic) setting without access to a consistent model of rewards and establish a minimax lower bound on the mean squa…

Cited by 254SourcePDFScholar
2016

Parallel and Distributed Block-Coordinate Frank-Wolfe Algorithms

ICML 2016poster

We study parallel and distributed Frank-Wolfe algorithms; the former on shared memory machines with mini-batching, and the latter in a delayed update framework. In both cases, we perform computations asynchronously whenever possible. We assume block-separable constraints as in Block-Coordinate Frank…

Cited by 56SourcePDFScholar
2016

Total Variation Classes Beyond 1d: Minimax Rates, and the Limitations of Linear Smoothers

NeurIPS 2016poster

We consider the problem of estimating a function defined over $n$ locations on a $d$-dimensional grid (having all side lengths equal to $n^{1/d}$). When the function is constrained to have discrete total variation bounded by $C_n$, we derive the minimax optimal (squared) $\ell_2$ estimation error r…

Cited by 92SourcePDFScholar
2015

A Deterministic Analysis of Noisy Sparse Subspace Clustering for Dimensionality-reduced Data

ICML 2015poster

Subspace clustering groups data into several lowrank subspaces. In this paper, we propose a theoretical framework to analyze a popular optimization-based algorithm, Sparse Subspace Clustering (SSC), when the data dimension is compressed via some random projection algorithms. We show SSC provably suc…

Cited by 42SourcePDFScholar
2015

Privacy for Free: Posterior Sampling and Stochastic Gradient Monte Carlo

ICML 2015poster

We consider the problem of Bayesian learning on sensitive datasets and present two simple but somewhat surprising results that connect Bayesian learning to “differential privacy”, a cryptographic approach to protect individual-level privacy while permitting database-level utility. Specifically, we s…

Cited by 304SourcePDFScholar