arXiv subjects

Jingzhao Zhang

Publications and source records attributed to Jingzhao Zhang.

At least 19 recordsLinked to original sources

arXiv2026

Optimal High-Order Methods for Solving Monotone Variational Inequalities

We study second- and higher-order methods for solving smooth monotone variational inequalities (MVI). Monteiro and Svaiter (SIAM J. Optim., 2012) showed that a second-order method, NPE, converges at a rate of $\mathcal{O}(T^{-1.5})$. For convex-concave minimax optimization, a subclass of MVI problems, Chen, Liu, Luo, and Zhang (COLT 2025) recently improved this rate to $\tilde{\mathcal{O}}( T^{-1.75})$. However, the result has a substantial gap compared to the lower bound of $Ω(T^{-2.5})$ established by Chen et al. (2026). In this paper, we propose a novel second-order method that achieves the optimal rate of $\mathcal{O}(T^{-2.5})$. Our algorithm also extends to higher-order methods: for any integer $ p \ge 1$, we obtain a $p$th-order method with a convergence rate of $\mathcal{O}(T^{-(3p-1)/2})$, matching the known lower bounds and therefore establishing optimal rates across all orders.

math.OC↗

arXiv2026

A Lower Bound for the Heavy-Ball Method on Smooth Convex Functions

Can the classical Heavy-Ball method, with arbitrary horizon-dependent parameters chosen in advance, achieve Nesterov's $O(T^{-2})$ last-iterate rate on every smooth convex objective? We provide a negative answer. For every horizon $T\ge2$ and every predetermined schedule with nonnegative step sizes and momenta in $[0,1)$, there exists a convex $1$-smooth objective, with initialization distance at most one and zero initial velocity, for which the last iterate of the Heavy-Ball method satisfies \[ f(x_T)-f^\star=Ω\!\left(\frac{1}{T^α\log T}\right), \qquad α=\frac{1+\sqrt5}{2}. \] Thus even fully nonstationary, horizon-dependent tuning cannot give the classical Heavy-Ball method a Nesterov-rate guarantee on the smooth convex class.

math.OC↗

arXiv2026

Scaling Domain Data Repetition in LLM Pretraining

As large language models scale, their training-token budgets must also increase to maintain an appropriate tokens-per-parameter ratio ($\mathrm{TPP}$). However, high-quality domain data is much harder to scale than general web data. As model size and the training-token budget increase, its fraction in the training mixture tends to decrease. Repeating the available high-quality data provides an effective way to counteract this dilution, but excessive repetition may lead to overfitting. We study this trade-off under practical LLM scaling, where the training-token budget grows proportionally with model size. For a fixed domain, we first find that, surprisingly at a fixed $\mathrm{TPP}$, the optimal repetition count mildly increases with model size. Across different domains, we find that the optimal repetition count is strongly negatively correlated with the final validation loss of a domain: domains with lower loss can generally benefit from more repetitions. In contrast, the amount of unique domain data is only weakly related to the optimal repetition count. These findings suggest that repetition counts tuned on smaller proxy models with the same $\mathrm{TPP}$ can provide a practical estimate for larger models.

cs.AI↗

arXiv2026

On the Condition Number Dependency in Bilevel Optimization

Bilevel optimization minimizes an objective function, defined by an upper-level problem whose feasible region is the solution of a lower-level problem. We study the oracle complexity of finding an $ε$-stationary point with first-order methods when the upper-level problem is nonconvex, and the lower-level problem is strongly convex. Recent works achieve a $\tilde{\mathcal{O}}(\bar κ_y^{7/2} ε^{-2})$ upper bound that is near-optimal in $ε$. In this work, we establish a new $Ω(κ_y^{5/2} ε^{-2})$ lower bound, where $κ_y \le \bar κ_y$ is the lower-level condition number. Our lower bound establishes the first provable gap {in terms of condition number dependency} between bilevel problems and minimax problems in this setup, and \textit{is tight up to logarithmic factors when the lower-level function is quadratic.} Our lower bounds can be extended to various settings. (1) For second-order and arbitrarily smooth problems, we show lower bounds of $Ω(κ_y^{9/4} ε^{-7/4})$ and $Ω(κ_y^{13/6} ε^{-5/3})$, respectively. (2) For convex--strongly-convex problems, we improve the previously best lower bound (Ji and Liang, JMLR 2022) from $Ω(κ_y /\sqrtε)$ to $Ω(κ_y^{3/2} / \sqrtε)$. (3) For stochastic nonconvex--strongly-convex problems, we also show the lower bounds of $Ω(κ_y^4 ε^{-4})$ and $Ω(κ_y^{9/2} ε^{-4})$ for stochastic Hessian-vector-product and stochastic first-order methods, respectively.

math.OC↗

arXiv2026

Halpern Iteration Achieves $\tilde{\mathcal{O}}(ε^{-1/p})$ $p$th-Order Oracle Complexity for Monotone Variational Inequalities

We study second- and higher-order methods for solving smooth monotone variational inequalities (MVI). Monteiro and Svaiter (SIAM J. Optim., 2012) showed that a second-order method, NPE, converges at the rate of $\mathcal{O}(T^{-1.5})$. For convex-concave minimax optimization, a subset of MVI problems, Chen, Liu, Luo, and Zhang (COLT 2025) recently improved the complexity to $\tilde{\mathcal{O}}( T^{-1.75})$ . However, it is open whether the conjectured complexity for MVI can be improved. In this paper, by using a large-step inexact Halpern iteration, we propose a novel Halpern-NPE method that achieves an even faster rate of $\tilde{\mathcal{O}}(T^{-2})$ for solving MVIs. We also provide the $p$th-order generalization of our method. We first introduce an Anchored Tensor Method (ATM) that achieves the rate of $\mathcal{O}(T^{-(p-1)})$, and then combine it with the Halpern iteration to achieve a faster convergence rate of $\tilde{\mathcal{O}}(T^{-p})$. This improves all prior results for $p \ge 2$ and matches the classical extragradient method for $p=1$.

math.OC↗

arXiv2026

Optimal Convex Optimization with Inexact Second-Order Oracles

In this paper, we present a novel second-order method called Accelerated Inexact Newton Extragradient (AINE) for convex optimization using $δ$-inexact Hessians. We show that AINE can find an $ε$-solution in the inexact second-order oracle (ISO) complexity of $\mathcal{O}( (δ/ε)^{1/2} + (L_2/ε)^{2/7} )$ when the Hessian is $L_2$-Lipschitz continuous, and a better complexity of $\mathcal{O}( (δ/ε)^{1/2} + (L_3/ε)^{1/5} )$ when the third-order derivative is $L_3$-Lipschitz continuous. Notably, each iteration of our method can be conducted in the same running time as matrix multiplication up to logarithmic factors. In addition, we also establish matching oracle complexity lower bounds for both setups, demonstrating the optimality of our methods.

math.OC↗

arXiv2026

Faster Newton Methods for Convex and Nonconvex Optimization in Gradient Complexity

Second-order optimization methods are computationally expensive for large-scale problems. Recently, Doikov, Chayti, and Jaggi (ICML 2023) proposed the LazyCRN method that reduces computation by studying the gradient complexity of second-order methods. Their method can achieve a gradient complexity of $\mathcal{O}( \bar d + \bar d^{1/2} ε^{-3/2})$ and $\mathcal{O}( \bar d + \bar d^{1/2} ε^{-1/2})$ for nonconvex and convex optimization, respectively, where $\bar d$ is the effective dimension and $ε$ is the target precision. Very recently, Adil, Bullins, Sidford, and Zhang (NeurIPS 2025) improved the gradient complexity to $\mathcal{O}( \bar d + \bar d^{1/3} ε^{-3/2} \ln^{18} ε^{-1})$ for nonconvex optimization. However, the tightness of these methods remains open. In this work, we propose new methods that achieve an improved complexity of $\mathcal{O}( \bar d + \bar d^{1/3} ε^{-3/2})$ and $\mathcal{O}( (\bar d + \bar d^{13/21} ε^{-2/7}) \ln \bar d)$ for nonconvex and convex optimization, respectively, improving best-known results for both setups.

math.OC↗

arXiv2026

On the Nonlinearity of Learning Rate Scaling for LLM Training

Learning-rate transfer can reduce the cost of training large language models: instead of sweeping learning rates at target scale, practitioners extrapolate from smaller runs. Existing approaches often assume that the optimal learning rate follows a log-linear scaling law in data scale and model size. We carefully examine and evaluate this scaling law. In our empirical study of GPT-2--style models from 22M to 707M parameters trained on 5B to 100B tokens, the optimal learning rate develops upward curvature at larger scales, leading to inaccurate extrapolation. We find that this curvature largely disappears when learning rates are replaced by effective learning rate (the step size in normalized weight space), and when data $D$ extrapolation is used instead of model size $N$ extrapolation. Next, we explain nonlinearity in scaling: weight-norm converges to equilibrium slower when optimal learning is small, requiring a larger step size to reduce the transient phase. Experiments with AdamH, which directly controls the effective learning rate, further support this explanation.

cs.LG↗

arXiv2026

Defending Against Malicious Finetuning by Scaling Train-time Adversarial Attacks

Current open-weight large language models (LLMs) are prone to malicious finetuning attacks, which could compromise the safety alignment of LLMs with only a few steps of supervised finetuning (SFT) on poisoned datasets. Existing alignment-stage defenses are primarily designed to defend against attacks that use parameter-efficient finetuning methods. However, they fail to defend against stronger attacks that use full-parameter finetuning. In this paper, we propose Patcher, a method inspired by adversarial training and bi-level optimization, to combat such attacks. Patcher strengthens the simulated attack by scaling up the optimization steps in the adversarial loop, thus forcing the defender to find model parameters that are insensitive to stronger attacks. Furthermore, we propose an efficient parallel algorithm to implement Patcher, decreasing the wall-clock time of training while preserving Patcher's performance. Extensive experiments show that Patcher substantially improves the model's robustness compared to vanilla SFT alignment, and transfers to diverse attack scenarios and model sizes. Code is available at https://github.com/haomingwen/patcher.

cs.CL↗

arXiv2026

Near-Optimal Nonconvex-Strongly-Convex Bilevel Optimization with Fully First-Order Oracles

In this work, we consider bilevel optimization when the lower-level problem is strongly convex. Recent works show that with a Hessian-vector product (HVP) oracle, one can provably find an $ε$-stationary point within ${\mathcal{O}}(ε^{-2})$ oracle calls. However, the HVP oracle may be inaccessible or expensive in practice. Kwon et al. (ICML 2023) addressed this issue by proposing a first-order method that can achieve the same goal at a slower rate of $\tilde{\mathcal{O}}(ε^{-3})$. In this paper, we incorporate a two-time-scale update to improve their method to achieve the near-optimal $\tilde {\mathcal{O}}(ε^{-2})$ first-order oracle complexity. Our analysis is highly extensible. In the stochastic setting, our algorithm can achieve the stochastic first-order oracle complexity of $\tilde {\mathcal{O}}(ε^{-4})$ and $\tilde {\mathcal{O}}(ε^{-6})$ when the stochastic noises are only in the upper-level objective and in both level objectives, respectively. When the objectives have higher-order smoothness conditions, our deterministic method can escape saddle points by injecting noise, and can be accelerated to achieve a faster rate of $\tilde {\mathcal{O}}(ε^{-1.75})$ using Nesterov's momentum.

math.OC↗

arXiv2026

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

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 context-dependent. In this work, we systematically study the role of data difficulty in fine-tuning from both empirical and theoretical perspectives, and find that there is no universally optimal difficulty level; rather, its effectiveness depends on the dataset size. We show that for a fixed data budget, there exists an optimal data difficulty for SFT, and that this optimal difficulty shifts toward harder data as the data budget increases. To explain this phenomenon, we conduct controlled synthetic experiments that reveal a simple underlying mechanism: the interplay between the (in-distribution) generalization gap and the extrapolation gap. We further support this mechanism through a theoretical analysis using PAC-Bayesian generalization bounds. Overall, our results clarify how data size and difficulty jointly affect the trade-off between generalization and extrapolation in SFT, providing guidance for difficulty-based data selection under certain model and data conditions.

cs.LG↗

arXiv2026

Data Mixing Can Induce Phase Transitions in Knowledge Acquisition

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 knowledge-dense datasets, unlike training exclusively on knowledge-dense data (arXiv:2404.05405), does not always follow a smooth scaling law but can exhibit phase transitions with respect to the mixing ratio and model size. Through controlled experiments on a synthetic biography dataset mixed with web-scraped data, we demonstrate that: (1) as we increase the model size to a critical value, the model suddenly transitions from memorizing very few to most of the biographies; (2) below a critical mixing ratio, the model memorizes almost nothing even with extensive training, but beyond this threshold, it rapidly memorizes more biographies. We attribute these phase transitions to a capacity allocation phenomenon: a model with bounded capacity must act like a knapsack problem solver to minimize the overall test loss, and the optimal allocation across datasets can change discontinuously as the model size or mixing ratio varies. We formalize this intuition in an information-theoretic framework and reveal that these phase transitions are predictable, with the critical mixing ratio following a power-law relationship with the model size. Our findings highlight a concrete case where a good mixing recipe for large models may not be optimal for small models, and vice versa.

cs.LG↗

arXiv2026

Capacity-Aware Mixture Law Enables Efficient LLM Data Optimization

A data mixture refers to how different data sources are combined to train large language models, and selecting an effective mixture is crucial for optimal downstream performance. Existing methods either conduct costly searches directly on the target model or rely on mixture scaling laws that fail to extrapolate well to large model sizes. We address these limitations by introducing a compute-efficient pipeline for data mixture scaling. First, we propose CAMEL, a capacity-aware mixture law that models validation loss with the nonlinear interplay between model size and mixture. We also introduce a loss-to-benchmark prediction law that estimates benchmark accuracy from validation loss, enabling end-to-end performance prediction for the target model. Next, we study how to allocate a fixed compute budget across model scales to fit the law and reduce prediction error. Finally, we apply our method to Mixture-of-Experts models with up to 7B-A150M parameters to fit the law, and verify the optimal mixture derived from the law by extrapolating to a 55B-A1.2B target model. Compared to prior methods, we reduce mixture optimization costs by 50\% and improves downstream benchmark performance by up to 3\%.

cs.LG↗

arXiv2026

On Finding Small Hyper-Gradients in Bilevel Optimization: Hardness Results and Improved Analysis

Bilevel optimization reveals the inner structure of otherwise oblique optimization problems, such as hyperparameter tuning, neural architecture search, and meta-learning. A common goal in bilevel optimization is to minimize a hyper-objective that implicitly depends on the solution set of the lower-level function. Although this hyper-objective approach is widely used, its theoretical properties have not been thoroughly investigated in cases where the lower-level functions lack strong convexity. In this work, we first provide hardness results to show that the goal of finding stationary points of the hyper-objective for nonconvex-convex bilevel optimization can be intractable for zero-respecting algorithms. Then we study a class of tractable nonconvex-nonconvex bilevel problems when the lower-level function satisfies the Polyak-Łojasiewicz (PL) condition. We show a simple first-order algorithm can achieve better complexity bounds of $\tilde{\mathcal{O}}(ε^{-2})$, $\tilde{\mathcal{O}}(ε^{-4})$ and $\tilde{\mathcal{O}}(ε^{-6})$ in the deterministic, partially stochastic, and fully stochastic setting respectively. The complexities in the first two cases are optimal up to logarithmic factors.

math.OC↗

arXiv2026

Solving Convex-Concave Problems with $\tilde{\mathcal{O}}(ε^{-4/(3p+1)})$ $p$th-Order Oracle Complexity

When the objective has Lipschitz continuous $p$th-order derivatives, it is known that convex-concave minimax problems can be solved with $\mathcal{O}(ε^{-2/(p+1)})$ $p$th-order oracle calls. This complexity upper bound was speculated to be optimal as it is achieved by a natural generalization of the optimal first-order method. In this work, we show an improved upper bound of $\tilde{\mathcal{O}}(ε^{-4/(3p+1)})$ by applying the Monteiro-Svaiter acceleration. We also establish a lower complexity bound of $Ω(ε^{-2/(3p-1)})$, suggesting a gap still exists for $p \ge 2$.

math.OC↗

arXiv2026

PiERN: Token-Level Routing for Integrating High-Precision Computation and Reasoning

Tasks on complex systems require high-precision numerical computation to support decisions, but current large language models (LLMs) cannot integrate such computations as an intrinsic and interpretable capability with existing architectures. Multi-agent approaches can leverage external experts, but inevitably introduce communication overhead and suffer from inefficiency caused by limited scalability. To this end, we propose Physically-isolated Experts Routing Network (PiERN), an architecture for integrating computation and reasoning. Instead of the tool-use workflows or function-calling, PiERN endogenously integrates computational capabilities into neural networks after separately training experts, a text-to-computation module, and a router. At inference, the router directs computation and reasoning at the token level, thereby enabling iterative alternation within a single chain of thought. We evaluate PiERN on representative linear and nonlinear computation-reasoning tasks against LLM finetuning and the multi-agent system approaches. Results show that the PiERN architecture achieves not only higher accuracy than directly finetuning LLMs but also significant improvements in response latency, token usage, and GPU energy consumption compared with mainstream multi-agent approaches. PiERN offers an efficient, interpretable, and scalable paradigm for interfacing language models with scientific systems.

cs.LG↗

arXiv2026

Differences in Text Generated by Diffusion and Autoregressive Language Models

Diffusion language models (DLMs) are promising alternatives to autoregressive language models (ARMs), yet the intrinsic differences in their generated text remain underexplored. We first find empirically that off-the-shelf DLMs exhibit lower $n$-gram entropy, higher semantic coherence, and higher semantic diversity. To understand the cause, we conduct controlled experiments that decouple the effects of training objectives and decoding algorithms. Results suggest that the DLM training objective contributes to the increases in semantic coherence and semantic diversity, but has a minor influence on entropy. These differences are primarily driven by the bidirectional context; other components in the training objective, such as input masking, label masking, and the weighting function, have a much weaker influence. Further, our experiments demonstrate that the reduction in entropy stems from DLMs' decoding algorithms, particularly confidence-based remasking strategies. We provide a theoretical understanding for this entropy reduction phenomenon. Together, our work uncovers key mechanisms underlying the differences between DLMs and ARMs in text generation, and informs future design of training objectives and decoding algorithms in DLMs.

cs.CL↗

arXiv2026

QuestA: Expanding Reasoning Capacity in LLMs via Question Augmentation

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 problems more effectively? To address this challenge, we propose a simple yet effective strategy via Question Augmentation: introduce partial solutions during training to reduce problem difficulty and provide more informative learning signals. Our method, QuestA, when applied during RL training on math reasoning tasks, not only improves pass@1 but also pass@k-particularly on problems where standard RL struggles to make progress. This enables continual improvement over strong open-source models such as DeepScaleR and OpenMath Nemotron, further enhancing their reasoning capabilities. We achieve new state-of-the-art results on math benchmarks using 1.5B-parameter models: 72.50% (+10.73%) on AIME24, 62.29% (+12.79%) on AIME25, and 41.67% (+10.11%) on HMMT25. Code, data and model are available at https://github.com/foreverlasting1202/QuestA.

cs.CL↗