Searcharxiv⌕ Search

arXiv subjects

Jiaxiang Li

Publications and source records attributed to Jiaxiang Li.

At least 19 recordsLinked to original sources
arXiv2026

Aligning Large Language Models with Human Feedback: Mathematical Foundations and Algorithm Design

This article provides an introduction to the mathematical foundations and algorithmic frameworks used to align Large Language Models (LLMs) with human intentions, preferences, and values. We discuss standard alignment techniques, such as fine-tuning (SFT), reinforcement learning with human feedback (RLHF), and direct preference optimization (DPO). We also explore the theoretical underpinnings of learning from human preferences, drawing connections to inverse reinforcement learning (IRL) and discrete choice models. We present state-of-the-art algorithms in a tutorial style, discuss their advantages and limitations, and offer insights into practical implementation. Our exposition is intended to serve as a comprehensive resource for researchers and practitioners, providing both a foundational understanding of alignment methodologies and a framework for developing more robust and scalable alignment techniques.

cs.LG↗
arXiv2026

ARES: a plane-wave density-functional-theory code in Rust

ARES (Ab initio atomic mateRial modEling Software at JLU) is a plane-wave density-functional-theory (DFT) code written entirely in Rust as a layered 22-crate workspace. It supports norm-conserving pseudopotentials and the projector augmented-wave (PAW) method with the local-density approximation (LDA) and Perdew-Burke-Ernzerhof (PBE) functionals, Γ-only and explicit k-point sampling, smearing and mixing, forces, stress, relaxation, band structures, densities of states (DOS) and dispersion corrections; execution is serial or MPI-distributed over k-point pools, band groups and plane-wave rows. The code's distinguishing engineering choice is its validation apparatus. Every numerical component is checked against Quantum ESPRESSO driven from the same physical and numerical inputs, and the comparison is organized as numbered fixtures kept in the repository whose reference data is regenerated by a script instead of being captured once. On six materials spanning a semiconductor, three metals and two two-dimensional systems, band energies within one electronvolt of the Fermi level agree with the reference to a median below 0.4 meV. The pseudopotentials the code reads are validated against all-electron reference data as well: the median delta factor is 0.696 meV/atom over all 65 paired elements. Restrictions are stated with the capabilities: two functionals, no spin-polarized or spin-orbit path, a stress implementation with an explicit scope, and no accelerator execution.

cond-mat.mtrl-sci↗
arXiv2026

Accurately Generalizing Deep-Learning Hamiltonian Methods to Plane-Wave Density Functional Theory

Deep neural networks that predict density functional theory (DFT) Hamiltonians from material structures offer a route to efficient electronic-structure calculations. However, extending localized atomic-orbital (AO) Hamiltonian learning to plane-wave (PW) calculations with hybrid functionals requires a compact AO representation that preserves the effects of nonlocal operators. Here we propose a projectability-guided construction of nonorthogonal AO Hamiltonians from a small, selected set of PW eigenstates and eigenvalues. High-projectability states are retained, and selected higher-energy states complete the AO subspace, avoiding an empirical shift and its arbitrary contribution to the learning targets. Because the construction uses PW eigenstates and eigenvalues, it incorporates nonlocal contributions present in the reference calculation without explicitly reconstructing their AO matrix elements. Benchmarks for MoS2 and sodium phases show accurate reproduction of PW bands with smaller AO bases than real-space reconstruction. The resulting PBE and HSE Hamiltonians can be learned by equivariant neural networks, connecting PW reference calculations with electronic-structure prediction. This approach broadens the range of PW DFT calculations available for AO Hamiltonian learning and enables compact, accurate training targets for developing more widely applicable electronic-structure models.

cond-mat.mtrl-sci↗
arXiv2026

Demystifying Manifold Constraints in LLM Pre-training

The recent success of matrix optimizers (e.g., Muon) suggests that specific normalization of momentum, such as orthogonalization and row-wise normalization, benefits both the stability and acceleration of LLM training. Consequently, several recent studies have suggested that weights should also be normalized, leading to a Riemannian optimization problem. While such constrained training frameworks demonstrate superior performance, the effects of explicitly constraining weights, and their interaction with existing stabilization mechanisms, remain less understood. To bridge this gap, we study manifold constrained training dynamics through activation scales, rotational dynamics, and the update-to-weight ratio. We propose a Riemannian spectral steepest descent optimizer called MACRO, alongside a radius selection principle to serve as our testbed. Our analysis and numerical experiments reveal that RMSNorm and manifold constraints serve overlapping roles, and that weight decay can be completely eliminated when manifold constraints are applied. By controlling the update-to-weight ratio, constrained training significantly alleviates update cancellation, empirically demonstrating that MACRO is robust to low-precision computation and competitive with existing algorithms for standard LLM pre-training.

cs.LG↗
arXiv2026

Centered Permutation Prefixes for SGD with Random Reshuffling: Sharp Rates, Hölder Geometry, and Composite Proximal Extensions

We study stochastic gradient descent with random reshuffling for finite sums \[ F(x)=\frac1n\sum_{i=1}^n f_i(x). \] For fresh reshuffling with a constant component stepsize, if each $f_i$ has an $L$-Lipschitz gradient and the average $F$ is $μ$-strongly convex with a Lipschitz-continuous Hessian, we prove the last-epoch rate \[ \mathbb E[F(y_K)-F(x_\star)] =\widetilde O\!\left(T^{-2}+n^2T^{-3}\right), \qquad T=nK, \] matching the known quadratic lower bound in its $(n,K)$-dependence. The components may be nonconvex, and no componentwise Hessian continuity or separate bounded-iterate assumption is required. More generally, a $ν$-Hölder-continuous average Hessian adds only $\widetilde O(n^{1+ν}T^{-2-2ν})$, so every $ν\ge 1/2$ preserves the quadratic rate. Under convex components, a decreasing-stepsize result removes the large-epoch requirement and recovers the same two-term scale once $nK$ exceeds the condition-number scale. We also analyze epoch-wise ProxRR for $\mathcal P=F+ψ$. Writing $x^\dagger$ for the composite minimizer and $β_\star=\|\nabla F(x^\dagger)\|$, we prove \[ \mathbb E\|y_K-x^\dagger\|^2 =\widetilde O\!\left( \frac{β_\star^2}{K^2} +T^{-2}+n^2T^{-3} +n^{1+ν}T^{-2-2ν} \right). \] For $ν\ge 1/2$, we show that the $β_\star^2/K^2$ splitting term is unavoidable and obtain a matching lower bound up to logarithms in the stated constant-stepsize regime.

math.OC↗
arXiv2026

Towards joint scaling laws with optimal batch size schedules

Modern deep learning typically keeps the batch size static throughout training, thus overlooking the joint effect of learning rate and batch size on the training dynamics. In this paper, we study the deep learning dynamics through the lens of convex optimization and derive a joint characterization of loss in terms of both schedules, applicable to general optimizers and model architectures. This characterization yields a closed-form optimal batch size schedule for any prescribed learning rate schedule, and further leads to joint scaling laws that consistently outperform static batch size baselines, highlighting the significance of dynamic batch size schedule in large language model training.

cs.LG↗
arXiv2026

On Second-Order Methods for Bilevel Optimization

Bilevel optimization is an indispensable modeling tool for modern machine learning and engineering design. However, the theory and practice for finding second order stationary points in the context of bilevel optimization still remain largely unsettled. Even for bilevel optimization with strongly convex lower-level problem, the hyperfunction it induces is in general nonconvex. Although the Cubic Regularized Newton methods (CRN) famously achieve the optimal $\mathcal{O}(\varepsilon^{-1.5})$ SOSP (second-order stationary point) rate in single-level optimization, it is unclear how to control the accuracy of the hypergradient and hyper-Hessian computations in the context of applying the second-order methods to bilevel problems in order for the overall process to be efficient. In this paper, we set out to answer this question. In particular, we first formulate a double loop CRN baseline that achieves the optimal outer rate but requires repeated lower level solves. Next, we propose a single loop cubic regularized Newton algorithm that combines one lower-level gradient step with one Newton step for the hypergradient, and prove an overall deterministic $\mathcal{O}(\varepsilon^{-1.5})$ total oracle complexity, which is optimal. In addition, we illustrate that some intuitively simple modifications of our method may fail to hold up the convergence result. To the best of our knowledge, this is the first deterministic single loop method for unconstrained NCSC (non-convex upper-level and strongly convex lower-level) bilevel optimization setting that achieves the $\mathcal{O}(\varepsilon^{-1.5})$ optimal convergence rate for finding an $\varepsilon$-SOSP of the hyperfunction.

math.OC↗
arXiv2026

Why Muon Outperforms Adam: A Curvature Perspective

Muon improves training efficiency over Adam in large language-model training by about two times, but the local geometric source of this advantage remains unclear. Our work takes a first step toward demystifying Muon's superiority over Adam from a curvature perspective. First, we apply a second-order Taylor approximation to the training landscape and show that Muon achieves a larger one-step loss decrease than Adam at matched validation loss. The two optimizers have comparable first-order gains, but Muon consistently incurs a smaller second-order curvature penalty. Second, we decompose this curvature penalty into the squared update norm and Normalized Directional Sharpness (NDS). We find that Muon and Adam have comparable update norms, so Muon's smaller curvature penalty is driven by lower NDS, not update scale. Third, we study how training data and model structure shape Muon's NDS advantage. Using Zipf-Probabilistic Context-Free Grammar (PCFG) data with controlled imbalance, we show that data imbalance amplifies Muon's NDS advantage over Adam. A within-/cross-layer decomposition further shows that, in the middle and late stages of training, Muon's lower NDS is mainly sustained by smaller within-layer curvature. Beyond empirical evidence, we analyze stylized quadratic problems with heterogeneous curvature and gradient alignment toward high-curvature modes. We prove that Muon attains a smaller average NDS than GD by balancing update energy across curvature groups; when curvature heterogeneity is sufficiently strong, this also yields lower local quadratic loss after the same number of steps.

cs.LG↗
arXiv2026

TPS-Drive: Task-Guided Representation Purification for VLM-based Autonomous Driving

Vision-Language Models (VLMs) provide a promising foundation for autonomous driving planning, yet bridging semantic reasoning and precise 3D spatial forecasting remains a critical challenge. Existing representation strategies generally follow two paths: text-aligned methods flatten continuous spatial states into symbols, which compromises geometric structure and induces "spatial hallucinations"; dense visual methods preserve spatial topology but overwhelm standard tokenizers with redundant background textures, leading to "representation interference". To address these limitations, we introduce TPS-Drive, a novel framework centered on Task-Guided Representation Purification that empowers VLMs to Think in Purified Space. At its core, an Agent-Centric Tokenizer utilizes a task-guided vector quantization mechanism supervised by a frozen 3D detection head, which explicitly reallocates limited codebook capacity from pervasive static backgrounds to critical dynamic agents and effectively isolates spatial redundancy. Leveraging this purified spatial vocabulary, TPS-Drive employs a decoupled reasoning pipeline that sequentially performs scene understanding, future forecasting, and action generation. The framework is optimized via a progressive three-stage training paradigm, culminating in reward-driven refinement that surpasses pure imitation learning. Extensive experiments validate our approach: TPS-Drive achieves accurate agent spatial state forecasting and reduces collision rates in open-loop nuScenes evaluations, while establishing new safety records on the rigorous closed-loop NAVSIMv1 and NAVSIMv2 benchmarks.

cs.RO↗
arXiv2026

ConFi-GS Confidence-Guided High-Frequency Injection for 3D Gaussian Splatting Super-Resolution

Reconstructing high-quality 3D scenes from low-resolution multi-view images remains challenging for 3D Gaussian Splatting (3DGS), because insufficient high-frequency observations often lead to blurred textures, weak boundaries, and view-inconsistent details. Existing approaches either apply super-resolution guidance uniformly or localize enhancement regions based mainly on geometric sampling. However, they typically do not distinguish between two fundamentally different questions: where additional detail is needed, and whether the corresponding candidate high-frequency content is reliable enough to be internalized into a multi-view consistent 3D representation. In this paper, we propose a reliability-aware frequency modeling framework for low-resolution 3DGS reconstruction. The framework first estimates a geometry-guided detail-demand prior to locate regions that are likely under-detailed under low-resolution supervision. It then computes a frequency-aware reliability map to determine whether candidate high-frequency details are structurally supported, spectrally unresolved, and cross-view stable. Combining these signals yields a detail-injection map that guides where super-resolved details should be introduced during optimization. Based on this map, we design a unified optimization scheme comprising spatially selective supervision, coarse-to-fine frequency regularization, and reliability-aware Gaussian densification. This scheme controls where reliable details are injected, when high-frequency supervision is activated, and how unresolved yet reliable details are internalized into the Gaussian representation. Experiments on multiple benchmarks show improved fidelity and perceptual quality while suppressing unstable or view-inconsistent details.

cs.CV↗
arXiv2026

Memory-Efficient LLM Pretraining via Minimalist Optimizer Design

Training large language models (LLMs) relies on adaptive optimizers such as Adam, which introduce extra operations and require significantly more memory to maintain first- and second-order moments than SGD. While recent works such as GaLore, Fira and APOLLO have proposed state-compressed memory-efficient variants, a fundamental question remains: What are the minimum modifications to plain SGD needed to match state-of-the-art pretraining performance? We systematically investigate this question using a bottom-up approach, and identify two simple yet highly (memory- and compute-) efficient techniques: (1) column-wise gradient normalization (normalizing the gradient along the output dimension), that boosts SGD performance without momentum; and (2) applying first-order momentum only to the output layer, where gradient variance is highest. Combining these two techniques lead to SCALE (Stochastic Column-normAlized Last-layer momEntum), a simple optimizer for memory efficient pretraining. Across multiple models (60M-1B), SCALE matches or exceeds the performance of Adam while using only 35-45% of the total memory. It also consistently outperforms memory-efficient optimizers such as GaLore, Fira and APOLLO, making it a strong candidate for large-scale pretraining under memory constraints. For LLaMA 7B, SCALE outperforms the state-of-the-art memory-efficient methods APOLLO and Muon in both perplexity and memory consumption.

cs.LG↗
arXiv2026

A Tale of Two Problems: Multi-Task Bilevel Learning Meets Equality Constrained Multi-Objective Optimization

In recent years, bilevel optimization (BLO) has attracted significant attention for its broad applications in machine learning. However, most existing works on BLO remain confined to the single-task setting and rely on the lower-level strong convexity assumption, which significantly restricts their applicability to modern machine learning problems of growing complexity. In this paper, we make the first attempt to extend BLO to the multi-task setting under a relaxed lower-level general convexity (LLGC) assumption. To this end, we reformulate the multi-task bilevel learning (MTBL) problem with LLGC into an equality constrained multi-objective optimization (ECMO) problem. However, ECMO itself is a new problem that has not yet been studied in the literature. To address this gap, we first establish a new Karush-Kuhn-Tucker (KKT)-based Pareto stationarity as the convergence criterion for ECMO algorithm design. Based on this foundation, we propose a weighted Chebyshev (WC)-penalty algorithm that achieves a finite-time convergence rate of $O(ST^{-\frac{1}{2})$ to KKT-based Pareto stationarity in both deterministic and stochastic settings, where $S$ denotes the number of objectives, and $T$ is the total iterations. Moreover, by varying the preference vector over the $S$-dimensional simplex, our WC-penalty method systematically explores the Pareto front. Finally, solutions to the ECMO problem translate directly into solutions for the original MTBL problem, thereby closing the loop between these two foundational optimization frameworks.

cs.LG↗
arXiv2026

Temporal Sampling Frequency Matters: A Capacity-Aware Study of End-to-End Driving Trajectory Prediction

End to end (E2E) autonomous driving trajectory prediction is often trained with camera frames sampled at the highest available temporal frequency, assuming that denser sampling improves performance. We question this assumption by treating temporal sampling frequency as an explicit training set design variable. Starting from high frequency E2E driving datasets, we construct frequency sweep training sets by temporally subsampling camera frames along each trajectory. For each model dataset pair, we train and evaluate the same model under a fixed protocol, so the frequency response reflects how prediction performance changes with sampling frequency. We analyze this response from a capacity aware perspective. Sparse sampling may miss driving relevant cues, while dense sampling may add redundant visual content and off manifold noise. For finite capacity models, this can create a driving irrelevant capacity burden. We evaluate three smaller E2E models and a larger VLA style AutoVLA model on Waymo, nuScenes, and PAVE. Results show model and dataset dependent frequency responses. Smaller E2E models often show non monotonic or near plateau trends and achieve their best 3 second ADE at lower or intermediate frequencies. In contrast, AutoVLA achieves its best 3 second ADE and FDE at the highest evaluated frequency on all three datasets. Iteration matched controls suggest that the advantage of lower or intermediate frequencies for smaller models is not explained only by unequal training update counts. These findings show that temporal sampling frequency should be reported and tuned, rather than fixed to the highest available value.

cs.CV↗
arXiv2026

A Correspondence-Driven Approach for Bilevel Decision-making with Nonconvex Lower-Level Problems

We consider bilevel optimization problems with general nonconvex lower-level objectives and show that the classical hyperfunction-based formulation is unsettled, since the global minimizer of the lower-level problem is generally unattainable. To address this issue, we propose a correspondence-driven hyperfunction $ϕ^{\text{cd}}$. In this formulation, the follower is modeled not as a rational agent always attaining a global minimizer, but as an algorithm-based bounded rational agent whose decisions are produced by a fixed algorithm with initialization and step size. Since $ϕ^{\text{cd}}$ is generally discontinuous, we apply Gaussian smoothing to obtain a smooth approximation $ϕ^{\text{cd}}_ξ$, then show that its value and gradient converge to those of $ϕ^{\text{cd}}$. In the nonconvex setting, we identify that bifurcation phenomena, which arise when $g(x,\cdot)$ has a degenerate stationary point, pose a key challenge for hyperfunction-based methods. This is especially the case when $ϕ^{\text{cd}}_ξ$ is solved using gradient methods. To overcome this challenge, we analyze the geometric structure of the bifurcation set under some weak assumptions. Building on these results, we design a biased projected SGD-based algorithm SCiNBiO to solve $ϕ^{\text{cd}}_ξ$ with a cubic-regularized Newton lower-level solver. We also provide convergence guarantees and oracle complexity bounds for the upper level. Finally, we connect bifurcation theory from dynamical systems to the bilevel setting and define the notion of fold bifurcation points in this setting. Under the assumption that all degenerate stationary points are fold bifurcation points, we establish the oracle complexity of SCiNBiO for the lower-level problem.

math.OC↗
arXiv2026

DTP-Attack: A decision-based black-box adversarial attack on trajectory prediction

Trajectory prediction systems are critical for autonomous vehicle safety, yet remain vulnerable to adversarial attacks that can cause catastrophic traffic behavior misinterpretations. Existing attack methods require white-box access with gradient information and rely on rigid physical constraints, limiting real-world applicability. We propose DTP-Attack, a decision-based black-box adversarial attack framework tailored for trajectory prediction systems. Our method operates exclusively on binary decision outputs without requiring model internals or gradients, making it practical for real-world scenarios. DTP-Attack employs a novel boundary walking algorithm that navigates adversarial regions without fixed constraints, naturally maintaining trajectory realism through proximity preservation. Unlike existing approaches, our method supports both intention misclassification attacks and prediction accuracy degradation. Extensive evaluation on nuScenes and Apolloscape datasets across state-of-the-art models including Trajectron++ and Grip++ demonstrates superior performance. DTP-Attack achieves 41 - 81% attack success rates for intention misclassification attacks that manipulate perceived driving maneuvers with perturbations below 0.45 m, and increases prediction errors by 1.9 - 4.2 for accuracy degradation. Our method consistently outperforms existing black-box approaches while maintaining high controllability and reliability across diverse scenarios. These results reveal fundamental vulnerabilities in current trajectory prediction systems, highlighting urgent needs for robust defenses in safety-critical autonomous driving applications.

cs.RO↗
arXiv2026

The Impact of Ionic Anharmonicity on Superconductivity in Metal-Stuffed B-C Clathrates

Metal-stuffed B$-$C compounds with sodalite clathrate structure have captured increasing attention due to their predicted exceptional superconductivity above liquid nitrogen temperature at ambient pressure. However, by neglecting the quantum lattice anharmonicity, the existing studies may result in an incomplete understanding of such a lightweight system. Here, using state-of-the-art ab initio methods incorporating quantum effects and machine learning potentials, we revisit the properties of a series of $XY$$\text{B}_{6}\text{C}_{6}$ clathrates where $X$ and $Y$ are metals. Our findings show that ionic quantum and anharmonic effects can harden the $E_g$ and $E_u$ vibrational modes, enabling the dynamical stability of 15 materials previously considered unstable in the harmonic approximation, including materials with previously unreported ($XY$)$^{1+}$ state, which is demonstrated here to be crucial to reach high critical temperatures. Further calculations based on the anisotropic Migdal-Eliashberg equation demonstrate that the $T_\text{c}$ values for KRb$\text{B}_{6}\text{C}_{6}$ and Rb$\text{B}_{3}\text{C}_{3}$ among these stabilized compounds are 102 and 115 K at 0 and 15 GPa, respectively, both being higher than $T_\text{c}$ of 92 K of KPb$\text{B}_{6}\text{C}_{6}$ at the anharmonic level. These record-high $T_\text{c}$ values, surpassing liquid nitrogen temperatures, emphasize the importance of anharmonic effects in stabilizing B-C clathrates with large electron-phonon coupling strength and advancing the search for high-$T_\text{c}$ superconductivity at (near) ambient pressure.

cond-mat.supr-con↗
arXiv2025

A Framework for Quantifying How Pre-Training and Context Benefit In-Context Learning

Pre-trained large language models have demonstrated a strong ability to learn from context, known as in-context learning (ICL). Despite a surge of recent applications that leverage such capabilities, it is by no means clear, at least theoretically, how the ICL capabilities arise, and in particular, what is the precise role played by key factors such as pre-training procedure as well as context construction. In this work, we propose a new framework to analyze the ICL performance, for a class of realistic settings, which includes network architectures, data encoding, data generation, and prompt construction process. As a first step, we construct a simple example with a one-layer transformer, and show an interesting result, namely when the pre-train data distribution is different from the query task distribution, a properly constructed context can shift the output distribution towards the query task distribution, in a quantifiable manner, leading to accurate prediction on the query topic. We then extend the findings in the previous step to a more general case, and derive the precise relationship between ICL performance, context length and the KL divergence between pre-train and query task distribution. Finally, we provide experiments to validate our theoretical results.

cs.AI↗
arXiv2025

ADARL: Adaptive Low-Rank Structures for Robust Policy Learning under Uncertainty

Robust reinforcement learning (Robust RL) seeks to handle epistemic uncertainty in environment dynamics, but existing approaches often rely on nested min--max optimization, which is computationally expensive and yields overly conservative policies. We propose \textbf{Adaptive Rank Representation (AdaRL)}, a bi-level optimization framework that improves robustness by aligning policy complexity with the intrinsic dimension of the task. At the lower level, AdaRL performs policy optimization under fixed-rank constraints with dynamics sampled from a Wasserstein ball around a centroid model. At the upper level, it adaptively adjusts the rank to balance the bias--variance trade-off, projecting policy parameters onto a low-rank manifold. This design avoids solving adversarial worst-case dynamics while ensuring robustness without over-parameterization. Empirical results on MuJoCo continuous control benchmarks demonstrate that AdaRL not only consistently outperforms fixed-rank baselines (e.g., SAC) and state-of-the-art robust RL methods (e.g., RNAC, Parseval), but also converges toward the intrinsic rank of the underlying tasks. These results highlight that adaptive low-rank policy representations provide an efficient and principled alternative for robust RL under model uncertainty.

cs.LG↗