← Search

Jingfeng Wu

27 accepted papers

2026

Memory-Statistics Tradeoff in Continual Learning with Structural Regularization

ICLR 2026poster

We study the statistical performance of a continual learning problem with two linear regression tasks in a well-specified random design setting. We consider a structural regularization algorithm that incorporates a generalized $\ell_2$-regularization tailored to the Hessian of the previous task for…

Cited by 0SourceScholar
2026

Seesaw: Accelerating Training by Balancing Batch Size and Learning Rate Scheduling

ICLR 2026poster

Increasing the batch size during training --- a “batch ramp'' --- is a promising strategy to accelerate large language model pretraining. While for SGD, doubling the batch size can be equivalent to halving the learning rate, the optimal strategy for adaptive optimizers like Adam is less clear. As a…

Cited by 0SourceScholar
2025

Benefits of Early Stopping in Gradient Descent for Overparameterized Logistic Regression

ICML 2025poster

In overparameterized logistic regression, gradient descent (GD) iterates diverge in norm while converging in direction to the maximum $\ell_2$-margin solution---a phenomenon known as the implicit bias of GD. This work investigates additional regularization effects induced by early stopping in well-s…

Cited by 0SourcePDFScholar
2025

Gradient Descent Converges Arbitrarily Fast for Logistic Regression via Large and Adaptive Stepsizes

ICML 2025poster

We analyze the convergence of gradient descent (GD) with large, adaptive stepsizes for logistic regression on linearly separable data. The stepsize adapts to the current risk, scaled by a fixed base stepsize \eta. We prove that once the number of iterates t surpasses a margin-dependent threshold, th…

Cited by 0SourcePDFScholar
2025

How Does Critical Batch Size Scale in Pre-training?

ICLR 2025poster

Training large-scale models under given resources requires careful design of parallelism strategies. In particular, the efficiency notion of critical batch size (CBS), concerning the compromise between time and compute, marks the threshold beyond which greater data parallelism leads to diminishing r…

2025

Implicit Bias of Gradient Descent for Non-Homogeneous Deep Networks

ICML 2025poster

We establish the asymptotic implicit bias of gradient descent (GD) for generic non-homogeneous deep networks under exponential loss. Specifically, we characterize three key properties of GD iterates starting from a sufficiently small empirical risk, where the threshold is determined by a measure of…

Cited by 0SourcePDFScholar
2025

Large Stepsizes Accelerate Gradient Descent for Regularized Logistic Regression

NeurIPS 2025poster

We study *gradient descent* (GD) with a constant stepsize for $\ell_2$-regularized logistic regression with linearly separable data. Classical theory suggests small stepsizes to ensure monotonic reduction of the optimization objective, achieving exponential convergence in $\widetilde{\mathcal{O}}(\k…

Cited by 0SourceScholar
2024

How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?

ICLR 2024spotlight

Transformers pretrained on diverse tasks exhibit remarkable in-context learning (ICL) capabilities, enabling them to solve unseen tasks solely based on input contexts without adjusting model parameters. In this paper, we study ICL in one of its simplest setups: pretraining a single-layer linear atte…

Cited by 70SourcePDFScholar
2024

In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization

NeurIPS 2024poster

We study the \emph{in-context learning} (ICL) ability of a \emph{Linear Transformer Block} (LTB) that combines a linear attention component and a linear multi-layer perceptron (MLP) component. For ICL of linear regression with a Gaussian prior and a \emph{non-zero mean}, we show that LTB can achiev…

Cited by 16SourcePDFScholar
2024

Large Stepsize Gradient Descent for Non-Homogeneous Two-Layer Networks: Margin Improvement and Fast Optimization

NeurIPS 2024poster

The typical training of neural networks using large stepsize gradient descent (GD) under the logistic loss often involves two distinct phases, where the empirical risk oscillates in the first phase but decreases monotonically in the second phase. We investigate this phenomenon in two-layer networks…

Cited by 6SourcePDFScholar
2024

Risk Bounds of Accelerated SGD for Overparameterized Linear Regression

ICLR 2024poster

Accelerated stochastic gradient descent (ASGD) is a workhorse in deep learning and often achieves better generalization performance than SGD. However, existing optimization theory can only explain the faster convergence of ASGD, but cannot explain its better generalization. In this paper, we study t…

Cited by 3SourcePDFScholar
2024

Scaling Laws in Linear Regression: Compute, Parameters, and Data

NeurIPS 2024poster

Empirically, large-scale deep learning models often satisfy a neural scaling law: the test error of the trained model improves polynomially as the model size and data size grow. However, conventional wisdom suggests the test error consists of approximation, bias, and variance errors, where the varia…

Cited by 17SourcePDFScholar
2023

Finite-Sample Analysis of Learning High-Dimensional Single ReLU Neuron

ICML 2023poster

This paper considers the problem of learning single ReLU neuron with squared loss (a.k.a., ReLU regression) in the overparameterized regime, where the input dimension can exceed the number of samples. We analyze a Perceptron-type algorithm called GLM-tron [Kakade et al. 2011], and provide its dimens…

Cited by 7SourcePDFScholar
2023

Implicit Bias of Gradient Descent for Logistic Regression at the Edge of Stability

NeurIPS 2023spotlight

Recent research has observed that in machine learning optimization, gradient descent (GD) often operates at the edge of stability (EoS) [Cohen et al., 2021], where the stepsizes are set to be large, resulting in non-monotonic losses induced by the GD iterates. This paper studies the convergence and…

Cited by 25SourcePDFScholar
2023

Private Federated Frequency Estimation: Adapting to the Hardness of the Instance

NeurIPS 2023poster

In federated frequency estimation (FFE), multiple clients work together to estimate the frequency of their local data by communicating with a server, while maintaining the security constraint of $\mathtt{secsum}$ where the server can only access the sum of client-held vectors. For FFE with a single…

Cited by 0SourcePDFScholar
2022

Gap-Dependent Unsupervised Exploration for Reinforcement Learning

AISTATS 2022poster

For the problem of task-agnostic reinforcement learning (RL), an agent first collects samples from an unknown environment without the supervision of reward signals, then is revealed with a reward and is asked to compute a corresponding near-optimal policy. Existing approaches mainly concern the wors…

2022

Last Iterate Risk Bounds of SGD with Decaying Stepsize for Overparameterized Linear Regression

ICML 2022oral

Stochastic gradient descent (SGD) has been shown to generalize well in many deep learning applications. In practice, one often runs SGD with a geometrically decaying stepsize, i.e., a constant initial stepsize followed by multiple geometric stepsize decay, and uses the last iterate as the output. Th…

Cited by 32SourcePDFScholar
2022

Risk Bounds of Multi-Pass SGD for Least Squares in the Interpolation Regime

NeurIPS 2022accept

Stochastic gradient descent (SGD) has achieved great success due to its superior performance in both optimization and generalization. Most of existing generalization analyses are made for single-pass SGD, which is a less practical variant compared to the commonly-used multi-pass SGD. Besides, theore…

Cited by 10SourcePDFScholar
2022

The Power and Limitation of Pretraining-Finetuning for Linear Regression under Covariate Shift

NeurIPS 2022accept

We study linear regression under covariate shift, where the marginal distribution over the input covariates differs in the source and the target domains, while the conditional distribution of the output given the input covariates is similar across the two domains. We investigate a transfer learning…

Cited by 24SourcePDFScholar
2021

Accommodating Picky Customers: Regret Bound and Exploration Complexity for Multi-Objective Reinforcement Learning

NeurIPS 2021poster

In this paper we consider multi-objective reinforcement learning where the objectives are balanced using preferences. In practice, the preferences are often given in an adversarial manner, e.g., customers can be picky in many applications. We formalize this problem as an episodic learning problem on…

2021

Direction Matters: On the Implicit Bias of Stochastic Gradient Descent with Moderate Learning Rate

ICLR 2021poster

Understanding the algorithmic bias of stochastic gradient descent (SGD) is one of the key challenges in modern machine learning and deep learning theory. Most of the existing works, however, focus on very small or even infinitesimal learning rate regime, and fail to cover practical scenarios where t…

Cited by 47SourcePDFScholar
2021

The Benefits of Implicit Regularization from SGD in Least Squares Problems

NeurIPS 2021poster

Stochastic gradient descent (SGD) exhibits strong algorithmic regularization effects in practice, which has been hypothesized to play an important role in the generalization of modern machine learning approaches. In this work, we seek to understand these issues in the simpler setting of linear regre…

Cited by 40SourcePDFScholar
2020

Obtaining Adjustable Regularization for Free via Iterate Averaging

ICML 2020poster

Regularization for optimization is a crucial technique to avoid overfitting in machine learning. In order to obtain the best performance, we usually train a model by tuning the regularization parameters. It becomes costly, however, when a single round of training takes significant amount of time. Ve…

2020

On the Noisy Gradient Descent that Generalizes as SGD

ICML 2020poster

The gradient noise of SGD is considered to play a central role in the observed strong generalization abilities of deep learning. While past studies confirm that the magnitude and the covariance structure of gradient noise are critical for regularization, it remains unclear whether or not the class o…

2019

The Anisotropic Noise in Stochastic Gradient Descent: Its Behavior of Escaping from Sharp Minima and Regularization Effects

ICML 2019oral

Understanding the behavior of stochastic gradient descent (SGD) in the context of deep neural networks has raised lots of concerns recently. Along this line, we study a general form of gradient based optimization dynamics with unbiased noise, which unifies SGD and standard Langevin dynamics. Through…