← Search

Jingzhao Zhang

36 accepted papers

2026

Data Difficulty and the Generalization–Extrapolation Tradeoff in LLM Fine-Tuning

ICML 2026poster

Data selection during supervised fine-tuning (SFT) can critically change the behavior of large language models (LLMs). Although existing work has studied the effect of selecting data based on heuristics such as perplexity, difficulty, or length, the reported findings are often inconsistent or contex…

Cited by 0SourceScholar
2026

Faster Gradient Methods for Highly-smooth Stochastic Bilevel Optimization

ICLR 2026poster

This paper studies the complexity of finding an $\epsilon$-stationary point for stochastic bilevel optimization when the upper-level problem is nonconvex and the lower-level problem is strongly convex. Recent work proposed the first-order method, F${}^2$SA, achieving the $\tilde{\mathcal{O}}(\epsilo…

Cited by 0SourceScholar
2026

QuestA: Expanding Reasoning Capacity in LLMs via Question Augmentation

ICLR 2026poster

Reinforcement learning (RL) has emerged as a central paradigm for training large language models (LLMs) in reasoning tasks. Yet recent studies question RL’s ability to incentivize reasoning capacity beyond the base model. This raises a key challenge: how can RL be adapted to solve harder reasoning p…

Cited by 0SourcecodeScholar
2026

Scaling Knowledge Editing in LLMs to 100,000 Facts with Neural KV Database

ICLR 2026poster

Efficiently editing knowledge stored in Large Language Models (LLMs) enables model updates without large-scale training. One promising solution is Locate-and-Edit (L\&E), allowing simultaneous modifications of a massive number of factual knowledge. However, such editing may compromise the general ab…

Cited by 0SourcecodeScholar
2025

Data Mixing Can Induce Phase Transitions in Knowledge Acquisition

NeurIPS 2025spotlight

Large Language Models (LLMs) are typically trained on data mixtures: most data come from web scrapes, while a small portion is curated from high-quality sources with dense domain-specific knowledge. In this paper, we show that when training LLMs on such data mixtures, knowledge acquisition from know…

Cited by 0SourceScholar
2025

Finite Sample Analyses for Continuous-time Linear Systems: System Identification and Online Control

NeurIPS 2025poster

Real world evolves in continuous time but computations are done from finite samples. Therefore, we study algorithms using finite observations in continuous-time linear dynamical systems. We first study the system identification problem, and propose a first non-asymptotic error analysis with finite o…

Cited by 0SourceScholar
2025

From Sparse Dependence to Sparse Attention: Unveiling How Chain-of-Thought Enhances Transformer Sample Efficiency

ICLR 2025poster

Chain-of-thought (CoT) significantly enhances the reasoning performance of large language models (LLM). While current theoretical studies often attribute this improvement to increased expressiveness and computational capacity, we argue that expressiveness is not the primary limitation in the LLM re…

2025

Generalization Lower Bounds for GD and SGD in Smooth Stochastic Convex Optimization

AISTATS 2025poster

This work studies the generalization error of gradient methods. More specifically, we focus on how training steps $T$ and step-size $\eta$ might affect generalization in smooth stochastic convex optimization (SCO) problems. Recent works show that in some cases longer training can hurt generalization…

Cited by 0SourceScholar
2025

Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks?

ICML 2025poster

Large language models (LLMs) exhibit remarkable task generalization, solving tasks they were never explicitly trained on with only a few demonstrations. This raises a fundamental question: When can learning from a small set of tasks generalize to a large task family? In this paper, we investigate t…

Cited by 0SourcePDFScholar
2025

Towards Black-Box Membership Inference Attack for Diffusion Models

ICML 2025poster

Given the rising popularity of AI-generated art and the associated copyright concerns, identifying whether an artwork was used to train a diffusion model is an important research topic. The work approaches this problem from the membership inference attack (MIA) perspective. We first identify the lim…

Cited by 4SourcePDFScholar
2025

Understanding Nonlinear Implicit Bias via Region Counts in Input Space

ICML 2025poster

One explanation for the strong generalization ability of neural networks is implicit bias. Yet, the definition and mechanism of implicit bias in non-linear contexts remains little understood. In this work, we propose to characterize implicit bias by the count of connected regions in the input space…

Cited by 0SourcePDFScholar
2024

A Quadratic Synchronization Rule for Distributed Deep Learning

ICLR 2024poster

In distributed deep learning with data parallelism, synchronizing gradients at each training step can cause a huge communication overhead, especially when many nodes work together to train large models. Local gradient methods, such as Local SGD, address this issue by allowing workers to compute lo…

2024

Functionally Constrained Algorithm Solves Convex Simple Bilevel Problem

NeurIPS 2024poster

This paper studies simple bilevel problems, where a convex upper-level function is minimized over the optimal solutions of a convex lower-level problem. We first show the fundamental difficulty of simple bilevel problems, that the approximate optimal value of such problems is not obtainable by first…

Cited by 0SourcePDFScholar
2024

Online Control with Adversarial Disturbance for Continuous-time Linear Systems

NeurIPS 2024poster

We study online control for continuous-time linear systems with finite sampling rates, where the objective is to design an online procedure that learns under non-stochastic noise and performs comparably to a fixed optimal linear controller. We present a novel two-level online algorithm, by integrat…

Cited by 0SourcePDFScholar
2024

Online Policy Optimization for Robust Markov Decision Process

UAI 2024poster

Reinforcement learning (RL) has exceeded human performance in many synthetic settings such as video games and Go. However, real-world deployment of end-to-end RL models is less common, as RL models can be very sensitive to perturbations in the environment. The robust Markov decision process (MDP) fr…

2023

Benign Overfitting in Classification: Provably Counter Label Noise with Larger Models

ICLR 2023poster

Studies on benign overfitting provide insights for the success of overparameterized deep learning models. In this work, we examine whether overfitting is truly benign in real-world classification tasks. We start with the observation that a ResNet model overfits benignly on Cifar10 but not benignly o…

Cited by 8SourcePDFScholar
2023

Fast Conditional Mixing of MCMC Algorithms for Non-log-concave Distributions

NeurIPS 2023poster

MCMC algorithms offer empirically efficient tools for sampling from a target distribution $\pi(x) \propto \exp(-V(x))$. However, on the theory side, MCMC algorithms suffer from slow mixing rate when $\pi(x)$ is non-log-concave. Our work examines this gap and shows that when Poincar\'e-style inequali…

Cited by 7SourcePDFScholar
2023

Iteratively Learn Diverse Strategies with State Distance Information

NeurIPS 2023poster

In complex reinforcement learning (RL) problems, policies with similar rewards may have substantially different behaviors. It remains a fundamental challenge to optimize rewards while also discovering as many *diverse* strategies as possible, which can be crucial in many practical applications. Our…

Cited by 4SourcePDFScholar
2023

On the Overlooked Pitfalls of Weight Decay and How to Mitigate Them: A Gradient-Norm Perspective

NeurIPS 2023poster

Weight decay is a simple yet powerful regularization technique that has been very widely used in training of deep neural networks (DNNs). While weight decay has attracted much attention, previous studies fail to discover some overlooked pitfalls on large gradient norms resulted by weight decay. In t…

2022

Beyond Worst-Case Analysis in Stochastic Approximation: Moment Estimation Improves Instance Complexity

ICML 2022spotlight

We study oracle complexity of gradient based methods for stochastic approximation problems. Though in many settings optimal algorithms and tight lower bounds are known for such problems, these optimal algorithms do not achieve the best performance when used in practice. We address this theory-practi…

Cited by 1SourcePDFScholar
2022

Neural Network Weights Do Not Converge to Stationary Points: An Invariant Measure Perspective

ICML 2022spotlight

This work examines the deep disconnect between existing theoretical analyses of gradient-based algorithms and the practice of training deep neural networks. Specifically, we provide numerical evidence that in large-scale neural network training (e.g., ImageNet + ResNet101, and WT103 + TransformerXL…

2021

Complexity Lower Bounds for Nonconvex-Strongly-Concave Min-Max Optimization

NeurIPS 2021poster

We provide a first-order oracle complexity lower bound for finding stationary points of min-max optimization problems where the objective function is smooth, nonconvex in the minimization variable, and strongly concave in the maximization variable. We establish a lower bound of $\Omega\left(\sqrt{\k…

Cited by 52SourcePDFScholar
2021

Coping with Label Shift via Distributionally Robust Optimisation

ICLR 2021poster

The label shift problem refers to the supervised learning setting where the train and test label distributions do not match. Existing work addressing label shift usually assumes access to an unlabelled test sample. This sample may be used to estimate the test label distribution, and to then train a…

Cited by 91SourcePDFScholar
2021

Exposure Bias versus Self-Recovery: Are Distortions Really Incremental for Autoregressive Text Generation?

EMNLP 2021main

Exposure bias has been regarded as a central problem for auto-regressive language models (LM). It claims that teacher forcing would cause the test-time generation to be incrementally distorted due to the training-generation discrepancy. Although a lot of algorithms have been proposed to avoid teache…

2021

Fast Federated Learning in the Presence of Arbitrary Device Unavailability

NeurIPS 2021poster

Federated learning (FL) coordinates with numerous heterogeneous devices to collaboratively train a shared model while preserving user privacy. Despite its multiple advantages, FL faces new challenges. One challenge arises when devices drop out of the training process. In this case, the convergence o…

Cited by 114SourcePDFScholar
2021

Provably Efficient Algorithms for Multi-Objective Competitive RL

ICML 2021oral

We study multi-objective reinforcement learning (RL) where an agent’s reward is represented as a vector. In settings where an agent competes against opponents, its performance is measured by the distance of its average return vector to a target set. We develop statistically and computationally effic…

Cited by 28SourcePDFScholar
2020

Complexity of Finding Stationary Points of Nonconvex Nonsmooth Functions

ICML 2020poster

We provide the first non-asymptotic analysis for finding stationary points of nonsmooth, nonconvex functions. In particular, we study the class of Hadamard semi-differentiable functions, perhaps the largest class of nonsmooth functions for which the chain rule of calculus holds. This class contains…

2020

Why Gradient Clipping Accelerates Training: A Theoretical Justification for Adaptivity

ICLR 2020talk

We provide a theoretical explanation for the effectiveness of gradient clipping in training deep neural networks. The key ingredient is a new smoothness condition derived from practical neural network training examples. We observe that gradient smoothness, a concept central to the analysis of first-…

Cited by 609SourcecodeScholar
2020

Why are Adaptive Methods Good for Attention Models?

NeurIPS 2020poster

While stochastic gradient descent (SGD) is still the de facto algorithm in deep learning, adaptive methods like Clipped SGD/Adam have been observed to outperform SGD across important tasks, such as attention models. The settings under which SGD performs poorly in comparison to adaptive methods are n…

2018

Direct Runge-Kutta Discretization Achieves Acceleration

NeurIPS 2018spotlight

We study gradient-based optimization methods obtained by directly discretizing a second-order ordinary differential equation (ODE) related to the continuous limit of Nesterov's accelerated gradient method. When the function is smooth enough, we show that acceleration can be achieved by a stable disc…

Cited by 136SourcePDFScholar