arXiv subjects

Jingzhao Zhang

Publications and source records attributed to Jingzhao Zhang.

At least 19 recordsLinked to original sources

arXiv2026

Near-Optimal Second-Order Convex Finite-Sum Optimization under Individual Smoothness

We study the stochastic second-order oracle complexity of solving smooth convex finite-sum optimization with $n$ components. Very recently, Tsedenov \& Kovalev (2026) showed that the minimax optimal complexity under averaged smoothness and convexity is $\tildeΘ(n + n^{6/7} ε^{-2/7})$ by establishing nearly matching upper and lower bounds. In this paper, we show that individual smoothness and convexity provably enable further acceleration, where the lower complexity bound is $Ω(n+ n^{5/7} ε^{-2/7})$ under this stronger assumption. Then, by establishing a novel second-order self-bounding inequality, we design a new Taylor-corrected lazy Newton method, which can lead to a better complexity of $\tilde{\mathcal{O}}(n + n^{5/7} ε^{-2/7})$ by applying the optimal Monteiro-Svaiter acceleration, matching our lower bound up to only logarithmic factors.

math.OC↗

arXiv2026

Leveraging Dale's Principle as an Inductive Bias in Recurrent Neural Networks

One of the fundamental differences between biological neural networks and artificial neural networks (ANNs) is that biological neurons obey Dale's principle: each neuron projects either exclusively excitatory or exclusively inhibitory outputs. While this constraint is central to biological circuits, previous attempts to impose Dale's principle in ANNs have often led to training instability or degraded performance. In this work, we introduce EISep RNN, a recurrent neural network architecture that implements Dale's principle through a learnable excitatory-inhibitory separation. Across a series of learning settings, including multi-task learning, continual learning, and reinforcement learning, EISep achieves more stable training dynamics, competitive or improved performance, and sparser, more modular connectivity than vanilla RNNs and previous Dale-constrained models. EISep also exhibits stable and informative latent dynamics. These results establish Dale's principle as a biologically grounded and computationally effective inductive bias for recurrent neural networks.

cs.LG↗

arXiv2026

Beyond Data Schedule: Fast LLM Pretraining on Data Mixtures by Loss-Reweighting

Large language models often acquire specialized capabilities through successive pretraining stages, but these gains can come at the cost of forgetting earlier knowledge. In this work, we investigate whether such trade-offs can be improved by changing how strongly each domain contributes to learning while keeping its data exposure fixed. We find that, in controlled knowledge--mathematics experiments, loss reweighting produces a single checkpoint that surpasses the best knowledge and reasoning scores attained along either stage trajectory. This reveals an independent control that data-mixture ratios alone do not capture. Building on this finding, we propose LARE (Loss-Adaptive Reweighting under fixed Exposure), which automatically adjusts domain weights from running loss statistics. Motivated by a theoretical analysis of gradient noise, its normalized inverse-square-root rule balances relative loss scales without changing the sampled mixture or requiring additional runs to search for weights. Across knowledge, mathematics, and code, \Method{} improves all five evaluated benchmarks for a 1.5B model trained on 240B tokens, averaging +7.84 pp over a matched direct mixture, including gains of +17.06 pp on GSM8K and +9.00 pp on MATH. These findings demonstrate that adjusting domain loss scales can improve joint capabilities without additional domain exposure.

cs.LG↗

arXiv2026

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

Tasks on complex systems require high-precision numerical computation to support decisions. However, current large language models (LLMs), even with enhanced reasoning capabilities, cannot integrate such computations as an intrinsic and interpretable capability with existing architectures. To this end, we propose Physically-isolated Experts Routing Network (PiERN), an architecture that directs computation and reasoning at token level, thereby enabling iterative alternation within a single chain of thought. We systematically evaluate PiERN on representative computation-reasoning tasks, including PDEBench and battery management tasks. Results show that PiERN achieves not only higher accuracy than directly finetuning LLMs but also significant improvements in response latency, token usage, GPU energy consumption, and experts routing accuracy compared with mainstream multi-agent approaches, while exhibiting no significant degradation in performance on MMLU and GLUE benchmarks. PiERN offers an efficient, interpretable, and scalable paradigm for interfacing language models with scientific systems.

cs.LG↗

arXiv2026

Near-Optimal Convex Optimization with Lazy Second-Order Oracles

This paper studies the complexity of convex optimization using lazy second-order oracles (Doikov, Chayti, and Jaggi, ICML 2023), where an algorithm queries gradients every iteration and Hessians once per $m$ iterations. Under this setting, we show a lower bound of $Ω(m+ m^{1/7} ε^{-2/7})$ on the number of total iterations to find an $ε$-solution using a novel block zero-chain construction. Then we propose a novel method that achieves a new upper bound of $\tilde{\mathcal{O}}(m+ m^{1/7} ε^{-2/7})$, which significantly improves the prior one (Chen, Liu, Luo, and Zhang, COLT 2026) of $\tilde{\mathcal{O}}(m+ m^{13/21} ε^{-2/7})$ and is tight up to logarithmic factors.

math.OC↗

arXiv2026

Defending Against Malicious Finetuning by Scaling Train-time Adversarial Attacks

Malicious finetuning attacks pose a major safety threat against open-weight large language models (LLMs). However, existing alignment-stage defenses provide limited protection against strong attacks that use full-parameter finetuning. To address this issue, we propose Patcher, a novel and efficient adversarial training algorithm that alternates between an attack stage and a defense stage: the attack stage simulates many-step malicious finetuning and computes the resulting parameter displacement as an "attack vector", while the defense stage aims to preserve safe behavior under such perturbations. Compared with other adversarial training methods, a key feature of Patcher is that the attack vector can be reused across multiple defense updates, substantially reducing training costs. Theoretically, we characterize how the relative gradient estimation error is influenced by the number of attack steps during training and reuse-duration of attack vectors. Empirically, experiment results show that Patcher reduces Attack Success Rate by 67.5%, 53.4% and 54.8% on three benchmarks compared to the second-best method while preserving the model's utility. Moreover, Patcher remains effective under diverse model sizes and attack scenarios, and generalizes to prompt-based jailbreak attacks. Code is available at https://github.com/haomingwen/patcher

cs.CL↗

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

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↗