← Search

Niao He

70 accepted papers

2026

A Schrödinger Eigenfunction Method for Long-Horizon Stochastic Optimal Control

ICLR 2026poster

High-dimensional stochastic optimal control (SOC) becomes harder with longer planning horizons: existing methods scale linearly in the horizon $T$, with performance often deteriorating exponentially. We overcome these limitations for a subclass of linearly-solvable SOC problems—those whose uncontrol…

Cited by 0SourcecodeScholar
2026

Landing with the Score: Riemannian Optimization through Denoising

ICLR 2026poster

Under the \emph{data manifold hypothesis}, high-dimensional data concentrate near a low-dimensional manifold. We study Riemannian optimization when this manifold is only given implicitly through the data distribution, and standard geometric operations are unavailable. This formulation captures a br…

Cited by 0SourceScholar
2026

On the Benefits of Weight Normalization for Overparameterized Matrix Sensing

ICLR 2026poster

While normalization techniques are widely used in deep learning, their theoretical understanding remains relatively limited. In this work, we establish the benefits of (generalized) weight normalization (WN) applied to the overparameterized matrix sensing problem. We prove that WN with Riemannian op…

Cited by 0SourceScholar
2026

SALAAD: Sparse And Low-Rank Adaptation via ADMM for Large Language Model Inference

ICML 2026poster

Modern large language models are increasingly deployed under compute and memory constraints, making flexible control of model capacity a central challenge. While sparse and low-rank structures naturally trade off capacity and performance, existing approaches often rely on heuristic designs that igno…

Cited by 0SourceScholar
2026

When Scores Learn Geometry: Rate Separations under the Manifold Hypothesis

ICLR 2026poster

Score-based methods, such as diffusion models and Bayesian inverse problems, are often interpreted as learning the data distribution in the low-noise limit ($\sigma \to 0$). In this work, we propose an alternative perspective: their success arises from implicitly learning the data manifold rather th…

Cited by 0SourceScholar
2026

Zeroth-Order Optimization at the Edge of Stability

ICML 2026poster

Zeroth-order (ZO) methods are widely used when gradients are unavailable or prohibitively expensive, including black-box learning and memory-efficient fine-tuning of large models, yet their optimization dynamics in deep learning remain underexplored. In this work, we provide an explicit step size co…

Cited by 0SourceScholar
2025

AmorLIP: Efficient Language-Image Pretraining via Amortization

NeurIPS 2025poster

Contrastive Language-Image Pretraining (CLIP) has demonstrated strong zero-shot performance across diverse downstream text-image tasks. Existing CLIP methods typically optimize a contrastive objective using negative samples drawn from each minibatch. To achieve robust representation learning, these…

Cited by 0SourcecodeScholar
2025

Can RLHF be More Efficient with Imperfect Reward Models? A Policy Coverage Perspective

ICML 2025poster

Sample efficiency is critical for online Reinforcement Learning from Human Feedback (RLHF). While existing works investigate sample-efficient online exploration strategies, the potential of utilizing misspecified yet relevant reward models to accelerate learning remains underexplored. This paper stu…

2025

Efficiently Escaping Saddle Points for Policy Optimization

UAI 2025

Policy gradient (PG) is widely used in reinforcement learning due to its scalability and good performance. In recent years, several variance-reduced PG methods have been proposed with a theoretical guarantee of converging to an approximate first-order stationary point (FOSP) with the sample complexi

2025

Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-Tuning

NeurIPS 2025spotlight

Adapting large-scale foundational flow and diffusion generative models to optimize task-specific objectives while preserving prior information is crucial for real-world applications such as molecular design, protein docking, and creative image generation. Existing principled fine-tuning methods aim…

Cited by 0SourceScholar
2025

Learning to Steer Markovian Agents under Model Uncertainty

ICLR 2025poster

Designing incentives for an adapting population is a ubiquitous problem in a wide array of economic applications and beyond. In this work, we study how to design additional rewards to steer multi-agent systems towards desired policies \emph{without} prior knowledge of the agents' underlying learning…

2025

Provable Maximum Entropy Manifold Exploration via Diffusion Models

ICML 2025poster

Exploration is critical for solving real-world decision-making problems such as scientific discovery, where the objective is to generate truly novel designs rather than mimic existing data distributions. In this work, we address the challenge of leveraging the representational power of generative m…

Cited by 0SourcePDFScholar
2025

Scalable Neural Incentive Design with Parameterized Mean-Field Approximation

NeurIPS 2025poster

Designing incentives for a multi-agent system to induce a desirable Nash equilibrium is both a crucial and challenging problem appearing in many decision-making domains, especially for a large number of agents $N$. Under the exchangeability assumption, we formalize this incentive design (ID) problem…

Cited by 0SourceScholar
2025

Zeroth-Order Optimization Finds Flat Minima

NeurIPS 2025poster

Zeroth-order methods are extensively used in machine learning applications where gradients are infeasible or expensive to compute, such as black-box attacks, reinforcement learning, and language model fine-tuning. Existing optimization theory focuses on convergence to an arbitrary stationary point,…

Cited by 0SourceScholar
2024

Achieving Near-Optimal Convergence for Distributed Minimax Optimization with Adaptive Stepsizes

NeurIPS 2024poster

In this paper, we show that applying adaptive methods directly to distributed minimax problems can result in non-convergence due to inconsistency in locally computed adaptive stepsizes. To address this challenge, we propose D-AdaST, a Distributed Adaptive minimax method with Stepsize Tracking. The k…

Cited by 0SourcePDFScholar
2024

Automated Design of Affine Maximizer Mechanisms in Dynamic Settings

AAAI 2024technical

Dynamic mechanism design is a challenging extension to ordinary mechanism design in which the mechanism designer must make a sequence of decisions over time in the face of possibly untruthful reports of participating agents. Optimizing dynamic mechanisms for welfare is relatively well understood. Ho…

Cited by 9SourcePDFScholar
2024

DPZero: Private Fine-Tuning of Language Models without Backpropagation

ICML 2024poster

The widespread practice of fine-tuning large language models (LLMs) on domain-specific data faces two major challenges in memory and privacy. First, as the size of LLMs continues to grow, the memory demands of gradient-based training methods via backpropagation become prohibitively high. Second, giv…

2024

Generalization Bounds of Nonconvex-(Strongly)-Concave Stochastic Minimax Optimization

AISTATS 2024poster

This paper studies the generalization performance of algorithms for solving nonconvex-(strongly)-concave (NC-SC/NC-C) stochastic minimax optimization measured by the stationarity of primal functions. We first establish algorithm-agnostic generalization bounds via uniform convergence between the empi…

Cited by 5SourcePDFScholar
2024

Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems

NeurIPS 2024poster

Sharpness-aware minimization (SAM) improves generalization of various deep learning tasks. Motivated by popular architectures such as LoRA, we explore the implicit regularization of SAM for scale-invariant problems involving two groups of variables. Instead of focusing on commonly used sharpness, th…

Cited by 3SourcePDFScholar
2024

Model-Based RL for Mean-Field Games is not Statistically Harder than Single-Agent RL

ICML 2024poster

We study the sample complexity of reinforcement learning (RL) in Mean-Field Games (MFGs) with model-based function approximation that requires strategic exploration to find a Nash Equilibrium policy. We introduce the Partial Model-Based Eluder Dimension (P-MBED), a more effective notion to character…

2024

On the Statistical Efficiency of Mean-Field Reinforcement Learning with General Function Approximation

AISTATS 2024poster

In this paper, we study the fundamental statistical efficiency of Reinforcement Learning in Mean-Field Control (MFC) and Mean-Field Game (MFG) with general model-based function approximation. We introduce a new concept called Mean-Field Model-Based Eluder Dimension (MF-MBED), which characterizes the…

Cited by 16SourcePDFScholar
2024

Truly No-Regret Learning in Constrained MDPs

ICML 2024spotlight

Constrained Markov decision processes (CMDPs) are a common way to model safety constraints in reinforcement learning. State-of-the-art methods for efficiently solving CMDPs are based on primal-dual algorithms. For these algorithms, all currently known regret bounds allow for *error cancellations* --…

Cited by 12SourcePDFScholar
2023

Learning to Optimize with Stochastic Dominance Constraints

AISTATS 2023poster

In real-world decision-making, uncertainty is important yet difficult to handle. Stochastic dominance provides a theoretically sound approach to comparing uncertain quantities, but optimization with stochastic dominance constraints is often computationally expensive, which limits practical applicabi…

2023

On Imitation in Mean-field Games

NeurIPS 2023poster

We explore the problem of imitation learning (IL) in the context of mean-field games (MFGs), where the goal is to imitate the behavior of a population of agents following a Nash equilibrium policy according to some unknown payoff function. IL in MFGs presents new challenges compared to single-agent…

Cited by 2SourcePDFScholar
2023

Optimal Guarantees for Algorithmic Reproducibility and Gradient Complexity in Convex Optimization

NeurIPS 2023spotlight

Algorithmic reproducibility measures the deviation in outputs of machine learning algorithms upon minor changes in the training process. Previous work suggests that first-order methods would need to trade-off convergence rate (gradient complexity) for better reproducibility. In this work, we challen…

Cited by 5SourcePDFScholar
2023

Policy Mirror Ascent for Efficient and Independent Learning in Mean Field Games

ICML 2023poster

Mean-field games have been used as a theoretical tool to obtain an approximate Nash equilibrium for symmetric and anonymous $N$-player games. However, limiting applicability, existing theoretical results assume variations of a ``population generative model'', which allows arbitrary modifications of…

Cited by 35SourcePDFScholar
2023

Reinforcement Learning with General Utilities: Simpler Variance Reduction and Large State-Action Space

ICML 2023poster

We consider the reinforcement learning (RL) problem with general utilities which consists in maximizing a function of the state-action occupancy measure. Beyond the standard cumulative reward RL setting, this problem includes as particular cases constrained RL, pure exploration and learning from dem…

Cited by 19SourcePDFScholar
2023

Stochastic Policy Gradient Methods: Improved Sample Complexity for Fisher-non-degenerate Policies

ICML 2023poster

Recently, the impressive empirical success of policy gradient (PG) methods has catalyzed the development of their theoretical foundations. Despite the huge efforts directed at the design of efficient stochastic PG-type algorithms, the understanding of their convergence to a globally optimal policy i…

Cited by 50SourcePDFScholar
2023

Two Sides of One Coin: the Limits of Untuned SGD and the Power of Adaptive Methods

NeurIPS 2023poster

The classical analysis of Stochastic Gradient Descent (SGD) with polynomially decaying stepsize $\eta_t = \eta/\sqrt{t}$ relies on well-tuned $\eta$ depending on problem parameters such as Lipschitz smoothness constant, which is often unknown in practice. In this work, we prove that SGD with arbitr…

Cited by 30SourcePDFScholar
2022

A Natural Actor-Critic Framework for Zero-Sum Markov Games

ICML 2022spotlight

We introduce algorithms based on natural actor-critic and analyze their sample complexity for solving two player zero-sum Markov games in the tabular case. Our results improve the best-known sample complexities of policy gradient/actor-critic methods for convergence to Nash equilibrium in the multi-…

Cited by 31SourcePDFScholar
2022

Bring Your Own Algorithm for Optimal Differentially Private Stochastic Minimax Optimization

NeurIPS 2022accept

We study differentially private (DP) algorithms for smooth stochastic minimax optimization, with stochastic minimization as a byproduct. The holy grail of these settings is to guarantee the optimal trade-off between the privacy and the excess population loss, using an algorithm with a linear time-co…

Cited by 24SourcePDFScholar
2022

Faster Single-loop Algorithms for Minimax Optimization without Strong Concavity

AISTATS 2022poster

Gradient descent ascent (GDA), the simplest single-loop algorithm for nonconvex minimax optimization, is widely used in practical applications such as generative adversarial networks (GANs) and adversarial training. Albeit its desirable simplicity, recent work shows inferior convergence rates of GDA…

2022

Lifted Primal-Dual Method for Bilinearly Coupled Smooth Minimax Optimization

AISTATS 2022poster

We study the bilinearly coupled minimax problem: $\min_{x} \max_{y} f(x) + y^\top A x - h(y)$, where $f$ and $h$ are both strongly convex smooth functions and admit first-order gradient oracles. Surprisingly, no known first-order algorithms have hitherto achieved the lower complexity bound of $\Omeg…

Cited by 42SourcePDFScholar
2022

Nest Your Adaptive Algorithm for Parameter-Agnostic Nonconvex Minimax Optimization

NeurIPS 2022accept

Adaptive algorithms like AdaGrad and AMSGrad are successful in nonconvex optimization owing to their parameter-agnostic ability – requiring no a priori knowledge about problem-specific parameters nor tuning of learning rates. However, when it comes to nonconvex minimax optimization, direct extension…

Cited by 26SourcePDFScholar
2022

Sharp Analysis of Stochastic Optimization under Global Kurdyka-Lojasiewicz Inequality

NeurIPS 2022accept

We study the complexity of finding the global solution to stochastic nonconvex optimization when the objective function satisfies global Kurdyka-{\L}ojasiewicz (KL) inequality and the queries from stochastic gradient oracles satisfy mild expected smoothness assumption. We first introduce a general…

Cited by 32SourcePDFScholar
2022

Stochastic Second-Order Methods Improve Best-Known Sample Complexity of SGD for Gradient-Dominated Functions

NeurIPS 2022accept

We study the performance of Stochastic Cubic Regularized Newton (SCRN) on a class of functions satisfying gradient dominance property with $1\le\alpha\le2$ which holds in a wide range of applications in machine learning and signal processing. This condition ensures that any first-order stationary po…

Cited by 22SourcePDFScholar
2021

The complexity of nonconvex-strongly-concave minimax optimization

UAI 2021poster

This paper studies the complexity for finding approximate stationary points of nonconvex-strongly-concave (NC-SC) smooth minimax problems, in both general and averaged smooth finite-sum settings. We establish nontrivial lower complexity bounds for the two settings, respectively. Our result reveals s…

Cited by 85SourcePDFScholar
2020

Biased Stochastic First-Order Methods for Conditional Stochastic Optimization and Applications in Meta Learning

NeurIPS 2020poster

Conditional stochastic optimization covers a variety of applications ranging from invariant learning and causal inference to meta-learning. However, constructing unbiased gradient estimators for such problems is challenging due to the composition structure. As an alternative, we propose a biased sto…

Cited by 73SourcePDFScholar
2020

Global Convergence and Variance Reduction for a Class of Nonconvex-Nonconcave Minimax Problems

NeurIPS 2020poster

Nonconvex minimax problems appear frequently in emerging machine learning applications, such as generative adversarial networks and adversarial learning. Simple algorithms such as the gradient descent ascent (GDA) are the common practice for solving these nonconvex games and receive lots of empirica…

Cited by 126SourcePDFScholar
2020

The Devil is in the Detail: A Framework for Macroscopic Prediction via Microscopic Models

NeurIPS 2020spotlight

Macroscopic data aggregated from microscopic events are pervasive in machine learning, such as country-level COVID-19 infection statistics based on city-level data. Yet, many existing approaches for predicting macroscopic behavior only use aggregated data, leaving a large amount of fine-grained micr…

2019

Exponential Family Estimation via Adversarial Dynamics Embedding

NeurIPS 2019poster

We present an efficient algorithm for maximum likelihood estimation (MLE) of exponential family models, with a general parametrization of the energy function that includes neural networks. We exploit the primal-dual view of the MLE with a kinetics augmented model to obtain an estimate associated wi…

2019

Kernel Exponential Family Estimation via Doubly Dual Embedding

AISTATS 2019poster

We investigate penalized maximum log-likelihood estimation for exponential family distributions whose natural parameter resides in a reproducing kernel Hilbert space. Key to our approach is a novel technique, doubly dual embedding, that avoids computation of the partition function. This technique al…

2019

Learning Positive Functions with Pseudo Mirror Descent

NeurIPS 2019spotlight

The nonparametric learning of positive-valued functions appears widely in machine learning, especially in the context of estimating intensity functions of point processes. Yet, existing approaches either require computing expensive projections or semidefinite relaxations, or lack convexity and theor…

Cited by 11SourcePDFScholar
2018

Coupled Variational Bayes via Optimization Embedding

NeurIPS 2018poster

Variational inference plays a vital role in learning graphical models, especially on large-scale datasets. Much of its success depends on a proper choice of auxiliary distribution class for posterior approximation. However, how to pursue an auxiliary distribution class that achieves both good approx…

2018

SBEED: Convergent Reinforcement Learning with Nonlinear Function Approximation

ICML 2018oral

When function approximation is used, solving the Bellman optimality equation with stability guarantees has remained a major open problem in reinforcement learning for decades. The fundamental difficulty is that the Bellman operator may become an expansion in general, resulting in oscillating and eve…

Cited by 336SourcePDFScholar
2017

Learning from Conditional Distributions via Dual Embeddings

AISTATS 2017poster

Many machine learning tasks, such as learning with invariance and policy evaluation in reinforcement learning, can be characterized as problems of learning from conditional distributions. In such problems, each sample x itself is associated with a conditional distribution $p(z|x)$ represented by sam…

Cited by 156SourcePDFScholar
2015

Time-Sensitive Recommendation From Recurrent User Activities

NeurIPS 2015poster

By making personalized suggestions, a recommender system is playing a crucial role in improving the engagement of users in modern web-services. However, most recommendation algorithms do not explicitly take into account the temporal behavior and the recurrent activities of users. Two central but les…

Cited by 173SourcePDFScholar