arXiv subjects

Tiejun Li

Publications and source records attributed to Tiejun Li.

At least 19 recordsLinked to original sources

arXiv2026

CrystalREPA: Transferring Physical Priors from Universal MLIPs to Crystal Generative Models

Crystal generative models mainly learn what stable crystals look like, with little explicit supervision for what makes them stable. We reveal a substantial representation gap between state-of-the-art crystal generative models and pretrained universal machine learning interatomic potentials (MLIPs) via energy probing, and show this gap can be closed by a simple training-time alignment. We propose Crystal REPresentation Alignment (CrystalREPA), a plug-and-play framework that aligns the atom-wise hidden states of generative encoders with frozen MLIP representations through an element-aware contrastive objective, transferring stability-aware atomistic priors with marginal training overhead and no additional inference cost. Across three generative frameworks, ten MLIP teachers, and two benchmark datasets, CrystalREPA consistently improves the thermodynamic stability, structural validity, and structural fidelity of generated crystals. Equally important, we find that an MLIP's transfer effectiveness is poorly predicted by its accuracy on standard leaderboards (e.g., Matbench Discovery) but strongly predicted by the distinguishability of its atom-wise representation space, yielding a practical, accuracy-independent criterion for selecting MLIP teachers for generative transfer. Code is available at https://github.com/Chengqian-Zhang/CrystalREPA.

cond-mat.mtrl-sci↗

arXiv2026

SnakeCharmer: Orchestrating Utility, Timeliness, and Placement in GPU Prefetching

GPUs rely on massive thread-level parallelism to tolerate long memory latency, yet they still suffer severe stalls when many warps simultaneously wait for data from the memory system. Address-prediction-based GPU prefetchers can mitigate such stalls, but accurate prediction alone does not guarantee performance improvement: a predicted prefetch becomes beneficial only if it is worth issuing, arrives at the right time, and is retained without harmful interference. Because these requirements interact, optimizing any single dimension in isolation can simply move the bottleneck to another. We capture these requirements with a predicted-prefetch-to-realized-benefit perspective and propose SnakeCharmer, a unified backend prefetch-control framework for address-prediction-based GPU prefetchers. SnakeCharmer separates address prediction from prefetch realization: the frontend predictor generates prefetch candidates, while the backend control path jointly controls utility, timeliness, and placement by determining whether each candidate should be issued, when it should be issued, and where the returned data should be retained. SnakeCharmer realizes this control path with zero-value-aware filtering for utility, hierarchical feedback-driven distance tuning for timeliness, and a register-assisted prefetch buffer for placement. We design and implement SnakeCharmer primarily on top of Snake, and adapt the same backend-control principles to APOGEE to demonstrate that SnakeCharmer's benefits are not limited to a single frontend predictor. On an A100 GPU model, SnakeCharmer improves performance/energy efficiency by 12.8%/20.8% over Snake and by 1.7%/3.2% over APOGEE, with modest hardware overhead.

cs.AR↗

arXiv2026

NeuroPDE+: A Scalable Neuromorphic PDE Accelerator Based on Spintronic and Ferroelectric Devices

The pursuit of high-performance PDE solvers rests on three fundamental challenges: (i) the curse of dimensionality in kinetic and financial equations, (ii) the poor extrapolation of purely data-driven surrogates, and (iii) the widening gap between algorithm design and hardware specialization. To overcome these challenges, we present NeuroPDE+, a scalable neuromorphic PDE solver design based on spintronic and ferroelectric devices for accelerating PDE solutions. NeuroPDE+ consists of two dedicated units: a diffusion tracking unit (DTU), which emulates random walks on Markov chains through activations between hardware neurons, and a scattering tracking unit (STU), which samples non-local jumps via a multi-level probability tree. System-level simulations suggest that NeuroPDE+ achieves a squared error below 1e-2 in steady-state heat equation and particle transport problems. Simulation results further indicate that the DTU achieves up to a 315x performance gain over previous neuromorphic processors, and that the STU achieves a 1000x speedup compared to a general-purpose CPU. Co-designing algorithm and hardware with intrinsic stochasticity and non-volatile in-memory computing, NeuroPDE+ preliminarily explores a new paradigm for efficient and scalable neuromorphic PDE solvers. This approach could pave the way for probabilistic computing architectures in large-scale scientific simulations.

cs.AR↗

arXiv2026

Gradient-Update Mismatch: Rethinking Conflict-Free Training of Physics-Informed Neural Networks

Training Physics-Informed Neural Networks (PINNs) requires jointly optimizing physics residual and initial/boundary condition loss terms, which often induce conflicting gradients. Gradient surgery methods mitigate this issue by constructing directions from loss-specific gradients to reduce conflict before optimizer transformation. However, even when the constructed direction is conflict-free, this property may not be preserved after optimizer transformation. Let $a_t$ denote the direction constructed by gradient surgery, $u_t$ the optimizer proposal, and $\mathcal{C}_t$ the conflict-free cone induced by the loss-specific gradients. We show that modern optimizers can transform $a_t$ through mechanisms such as historical state, adaptive scaling, preconditioning, or decoupled weight decay, so $a_t \in \mathcal{C}_t$ does not generally imply $u_t \in \mathcal{C}_t$. We refer to this optimizer-induced discrepancy in conflict-freeness between $a_t$ and $u_t$ as Gradient-Update Mismatch (GUM). Accordingly, we propose Gradient-Update Alignment (GUA), which projects $u_t$ onto $\mathcal{C}_t$ to obtain the aligned update $p_t$ and applies $p_t$ to the parameters. When the optimizer maintains internal state, GUA further adjusts this state toward targets reconstructed from the applied update. We conduct extensive experiments and find that GUM is widespread across momentum, adaptive, and curvature-based optimizers, with conflict rates reaching up to 86.3%. Across all PINN settings, GUA achieves conflict-free applied updates and consistently improves various gradient surgery methods, reducing the relative $L_2$ error by up to 98.2% in individual settings. Data and code are available at https://github.com/JingXiao10/GUA.

cs.LG↗

arXiv2026

DDA-Net: Accurate TDD Channel Estimation via Deep Unfolding the Doppler-Delay-Angle Representation of Channel Signals

In TDD massive MIMO systems, channel estimation under sparse frequency-hopping pilots is challenging: each snapshot captures only one narrow pilot block that hops across frequency, with tens of milliseconds between adjacent snapshots. Finite-window leakage and off-grid effects weaken the ideal Doppler-delay-angle (DDA) sparsity, limiting both classical sparse recovery and purely data-driven approaches lacking an explicit structured transform-domain model. We propose DDA-Net, a model-driven 3D deep unfolding network for joint multi-snapshot channel state reconstruction. DDA-Net integrates an ADMM-based formulation with a closed-form data-consistency update that avoids tensor inversion, a lightweight Doppler-domain learned prior, and delay oversampling to mitigate basis mismatch. It consistently outperforms strong baselines across three channel settings from 3GPP TR 38.901: UMa-NLOS, UMi-NLOS, and CDL-B. Ablations confirm that window-level 3D processing and explicit Doppler modeling yield target-dependent improvements. With minimal target-domain fine-tuning, the UMa-pretrained model surpasses both its zero-shot version and counterparts trained from scratch with the same number of target-domain samples. The Doppler-domain design proves consistently superior to its time-domain equivalent, with a wider margin after fine-tuning. These results demonstrate that combining exact physical data consistency with a learned DDA-domain prior is an effective and sample-efficient approach to channel state acquisition under sparse frequency-hopping pilots.

eess.SP↗

arXiv2026

C2P-Cache: Scalable GPU L1 Cache Sharing via Concurrent Candidate Pruning

Modern GPUs rely on private per-SM L1 caches and a shared L2 cache, but this organization obscures cross-SM reuse: an L1 miss is typically forwarded to L2 even when the requested line already resides in a peer L1 cache, leading to redundant L2 access. Prior GPU L1-sharing designs attempt to recover such reuse through exact or broad remote-hit searches, which become increasingly difficult to scale and can interfere with the critical L1 miss path under high concurrency. %miss handling as more caches participate and more misses arrive concurrently. We observe that eliminating redundant L2 accesses does not require exact, chip-wide knowledge of private L1 contents. Instead, it requires only sufficient visibility to sharply narrow down a small set of candidate caches, leaving exact confirmation to a much smaller number of L1s. Based on this insight, we propose C2P-Cache, a scalable GPU L1-sharing mechanism that transforms remote-hit discovery from a chip-wide exact search problem into a lightweight filtering-and-confirmation process. C2P-Cache maintains compact Bloom-filter-based snapshots of private L1 tags, performs parallel chip-wide candidate filtering, and selectively probes only a small number of likely peer caches. To sustain high concurrency, C2P-Cache organizes filtering as bit-sliced matching over a banked and replicated snapshot matrix, enabling efficient, parallel processing of many concurrent misses without interfering with normal L1 accesses. Across a wide range of GPU workloads, C2P-Cache improves instructions per cycle (IPC) by up to 49.7\% and by 23.5\% on average for applications with high remote-L1 reuse and strong sensitivity to L2 latency, demonstrating that lightweight, scalable filtering can effectively unlock cross-SM reuse with modest overhead.

cs.AR↗

arXiv2026

iSMART: An Iterative Sampling-and-Regression Technique for Solving Martingale-Based PDEs

We propose the {\bf i}terative {\bf S}a{\bf M}pling-{\bf A}nd-{\bf R}egression {\bf T}echnique (iSMART) for high-dimensional martingale-based partial differential equations (PDEs) in this paper. By leveraging the $L^2$-projection property of conditional expectation and adopting the stop-gradient technique, iSMART reformulates the continuous martingale condition derived from PDEs into a sequence of tractable sampling-regression problems within an iterative framework. This approach relies solely on standard SDE path simulation and plain squared-error loss minimization, completely bypassing the need for adversarial optimization or nested expectation estimation in previous methods. iSMART accommodates linear, semi-linear, and fully nonlinear martingale-based PDEs within a unified iterative procedure. In particular, for fully nonlinear Hamilton-Jacobi-Bellman (HJB) equations, a freezing-and-compensating technique is introduced to strategically shift a portion of the nonlinearity into the SDE drift, thereby improving the convergence behavior of the iterations. Numerous numerical experiments on linear reaction-diffusion equations with sharp gradients, semilinear Burgers-type equations, and fully nonlinear HJB equations demonstrate the accuracy, efficiency, and robustness of the proposed approach in various high dimensions.

math.NA↗

arXiv2026

Iterative Neural Variational Solver of Committor Function in Under-damped Langevin Dynamics

The committor function describes the probability that a stochastic dynamical system, while transitioning between two metastable states, will first reach one state rather than the other. It is a central object in transition path theory, providing the accurate reaction rate and the characterization of the transition mechanism through reactive probability flux and the transition state ensemble. While the numerical methods for this function are well studied for over-damped Langevin dynamics where the friction damping coefficient goes to infinity, its computation in the general under-damped case remains an open challenge due to the lack of a variational structure in the underlying backward Kolmogorov equation. This restricts the quantitative investigation of the impact of the finite damping coefficient on rare transitions in real systems. In this work, by decomposing the infinitesimal generator of the underdamped Langevin dynamics into a self-adjoint and a skew-adjoint part, we develop an iterative scheme in which each subproblem has a variational form and can be solved efficiently using neural networks. Fundamentally, this approach can be viewed as a generalized Concus-Golub-Widlund (CGW) splitting method and each subproblem has a variational form, which can be solved efficiently. This iterative scheme also has the flexibility of introducing a preconditioning operator associated with a new overdamped Langevin dynamics, which naturally leads to the design of enhanced adaptive sampling for the subproblems. We validate the efficiency and accuracy of our method on benchmark systems including a one-dimensional double-well potential, the Muller-Brown potential, and alanine dipeptide. The proposed iterative variational framework provides a practical and extensible numerical tool for probing rare transition events in Langevin dynamics across varying damping coefficients.

physics.comp-ph↗

arXiv2026

FES-FM: Free Energy Surface Sampling via Reduced Flow Matching

Sampling the distribution of collective variables (CVs) and estimating the associated free energy surface are crucial problems in statistical physics, as they underpin a better understanding of chemical reactions and conformational transitions. Traditional methods usually rely on simulations in high-dimensional configuration space and project the resulting configurations onto the CV space. To improve sampling speed, we propose FES-FM, a reduced flow matching (FM) method for free energy surface (FES) sampling. We train a dynamical transport map in the CV space, thereby enabling direct sampling of CV distributions and reconstruction of the corresponding free energy surface. For many-particle systems, we construct a prior distribution based on the Hessian at a local minimum of the potential, which ensures both rotation-translation invariance and physically meaningful configurations. We evaluate the proposed method across a variety of potential functions and collective variables, including alanine dipeptide in implicit solvent as a molecular benchmark. Comparative experiments demonstrate that our approach significantly improves sampling speed while maintaining accuracy.

cs.LG↗

arXiv2026

ST-DDA: Dynamic Channel Estimation in the Doppler-Delay-Angle Domain via Sparse Subspace Tracking for TDD Systems

Accurate full-band channel acquisition in frequency-hopping time-division duplex (TDD) systems is challenging because each sounding slot observes only a limited frequency subband, while conventional single-slot recovery cannot fully exploit historical observations. We propose ST-DDA, an online sparse-subspace tracking framework for latest-slot reconstruction in the Doppler--delay--angle (DDA) domain. We first show that the Doppler-domain representation remains energy-concentrated under moderate channel variation, thereby supporting windowed DDA-domain sparse recovery. A local stability analysis further shows that the substantial overlap between adjacent windows enables the preceding-window estimate to warm-start each window-specific recovery problem, allowing the optimization progress to be carried across slots under a fixed per-slot iteration budget. For computational tractability, ST-DDA employs the alternating subspace method, which restricts the regularized least-squares fidelity updates to support-induced subspaces, together with position-encoded convolutional reweighting that exploits local angular and Doppler structures. Experiments show that reweighted ST-DDA achieves more accurate and reliable reconstruction than dynamic compressed-sensing baselines, particularly for longer sounding intervals and larger frequency-hopping periods, while maintaining comparable per-slot runtime.

eess.SP↗

arXiv2026

Riemannian Denoising Diffusion Probabilistic Models

We propose Riemannian Denoising Diffusion Probabilistic Models (RDDPMs) for learning distributions on submanifolds of Euclidean space that are level sets of functions, including most of the manifolds relevant to applications. Existing methods for generative modeling on manifolds rely on substantial geometric information such as geodesic curves or eigenfunctions of the Laplace-Beltrami operator and, as a result, they are limited to manifolds where such information is available. In contrast, our method, built on a projection scheme, can be applied to more general manifolds, as it only requires being able to evaluate the value and the first order derivatives of the function that defines the submanifold. We provide a theoretical analysis of our method in the continuous-time limit, which elucidates the connection between our RDDPMs and score-based generative models on manifolds. The capability of our method is demonstrated on datasets from previous studies and on new datasets sampled from two high-dimensional manifolds, i.e. $\mathrm{SO}(10)$ and the configuration space of molecular system alanine dipeptide with fixed dihedral angle.

cs.LG↗

arXiv2026

Full-field prediction for engineering-scale three-dimensional aircraft with multigrid-hierarchical learning

High-fidelity computational fluid dynamics is essential for aerospace design, but engineering-scale simulations of practical three-dimensional aircraft remain computationally expensive. Learning-based flow-field initialization can improve efficiency by reducing the numerical distance between the initial and converged solutions, yet existing deep learning approaches remain difficult to scale to large three-dimensional aircraft flows with multiscale regional heterogeneity. Most prior studies therefore focus on two-dimensional problems, surface quantities, integral aerodynamic coefficients, or simplified three-dimensional cases with limited grid resolution.Here we propose MHLF, a multigrid-hierarchical learning framework for accelerating engineering-scale aircraft flow simulations while preserving high-fidelity numerical accuracy. MHLF combines a topologically consistent geometric multigrid representation with a hierarchical strategy that captures regional flow heterogeneity during both prediction and subsequent CFD correction. Across three engineering-scale aircraft cases spanning Mach 0.15 to 6.0 and covering subsonic, transonic and supersonic regimes, MHLF accelerates convergence without sacrificing flow-field accuracy, achieving a 3 to 8 times efficiency improvement over conventional initialization. These results demonstrate practical full-flow-field prediction for large three-dimensional aircraft within the CFD domain and provide a foundation for data-driven acceleration of high-fidelity aircraft flow simulation.

physics.flu-dyn↗

arXiv2026

Covering-radius and Collinearity- Minimizing Pilots for Channel Estimation in TDD Systems

This letter studies pilot design for orthogonal frequency-division multiplexing-based time-division duplex (TDD) systems under a sliding-window latest-slot recovery framework that jointly exploits delay-Doppler sparsity across recent slots. Under contiguous-subband and fairness constraints, this viewpoint naturally leads to a geometry-aware time-frequency joint pilot assignment. We show that effective patterns should balance grid coverage and redundant-collinearity suppression, with an additional symmetry-avoidance refinement when complete collinearity elimination is infeasible. Based on these principles, we formulate a mixed-integer construction method compatible with practical TDD allocation. Numerical results show that minimum-coverage-radius and collinearity-control (MCC) pattern improves both surrogate geometry metrics and latest-slot recovery performance.

cs.IT↗

arXiv2026

WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport

The Wasserstein-Fisher-Rao (WFR) metric extends dynamic optimal transport (OT) by coupling displacement with change of mass, providing a principled geometry for modeling unbalanced snapshot dynamics. Existing WFR solvers, however, are often unstable, computationally expensive, and difficult to scale. Here we introduce WFR Flow Matching (WFR-FM), a simulation-free training algorithm that unifies flow matching with dynamic unbalanced OT. Unlike classical flow matching which regresses only a transport vector field, WFR-FM simultaneously regresses a vector field for displacement and a scalar growth rate function for birth-death dynamics, yielding continuous flows under the WFR geometry. Theoretically, we show that minimizing the WFR-FM loss exactly recovers WFR geodesics. Empirically, WFR-FM yields more accurate and robust trajectory inference in single-cell biology, reconstructing consistent dynamics with proliferation and apoptosis, estimating time-varying growth fields, and applying to generative dynamics under imbalanced data. It outperforms state-of-the-art baselines in efficiency, stability, and reconstruction accuracy. Overall, WFR-FM establishes a unified and efficient paradigm for learning dynamical systems from unbalanced snapshots, where not only states but also mass evolve over time. The Python code is available at https://github.com/QiangweiPeng/WFR-FM.

cs.LG↗

arXiv2026

Alternating Subspace Method for Sparse Recovery of Signals

This work has been submitted to the IEEE for possible publication. Copyright may be transferred without notice, after which this version may no longer be accessible. Numerous renowned algorithms for tackling the compressed sensing problem employ an alternating strategy, which typically involves data matching in one module and denoising in another. We present a novel approach, the Alternating Subspace Method (ASM), which integrates the principles of the greedy methods (e.g., the orthogonal matching pursuit type methods) and the splitting methods (e.g., the approximate message passing type methods). Crucially, ASM enhances the splitting method by achieving fidelity in a subspace-restricted fashion. \textcolor{black}{We reveal that such a restriction strategy guarantees global convergence via proximal residual control and establish its local geometric convergence on the LASSO problem.} Numerical experiments on the LASSO, channel estimation, and dynamic compressed sensing problems demonstrate its high convergence rate and its capacity to incorporate different prior distributions. Overall, the proposed method is promising in terms of efficiency, accuracy, and flexibility, and has the potential to be competitive in different sparse recovery applications.

cs.IT↗

arXiv2026

Prior-Guided Symbolic Regression: Towards Scientific Consistency in Equation Discovery

Symbolic Regression (SR) aims to discover interpretable equations from observational data, with the potential to reveal underlying principles behind natural phenomena. However, existing approaches often fall into the Pseudo-Equation Trap: producing equations that fit observations well but remain inconsistent with fundamental scientific principles. A key reason is that these approaches are dominated by empirical risk minimization, lacking explicit constraints to ensure scientific consistency. To bridge this gap, we propose PG-SR, a prior-guided SR framework built upon a three-stage pipeline consisting of warm-up, evolution, and refinement. Throughout the pipeline, PG-SR introduces a prior constraint checker that explicitly encodes domain priors as executable constraint programs, and employs a Prior Annealing Constrained Evaluation (PACE) mechanism during the evolution stage to progressively steer discovery toward scientifically consistent regions. Theoretically, we prove that PG-SR reduces the Rademacher complexity of the hypothesis space, yielding tighter generalization bounds and establishing a guarantee against pseudo-equations. Experimentally, PG-SR outperforms state-of-the-art baselines across diverse domains, maintaining robustness to varying prior quality, noisy data, and data scarcity.

cs.LG↗

arXiv2026

WFR-MFM: One-Step Inference for Dynamic Unbalanced Optimal Transport

Reconstructing dynamical evolution from limited observations is a fundamental challenge in single-cell biology, where dynamic unbalanced optimal transport provides a principled framework for modeling coupled transport and mass variation. However, existing approaches rely on trajectory simulation at inference time, making inference a key bottleneck for scalable applications. In this work, we propose a mean-flow framework for unbalanced flow matching that summarizes both transport and mass-growth dynamics over arbitrary time intervals using mean velocity and mass-growth fields, enabling fast one-step generation without trajectory simulation. To solve dynamic unbalanced optimal transport under the Wasserstein-Fisher-Rao geometry, we further build on this framework to develop Wasserstein-Fisher-Rao Mean Flow Matching (WFR-MFM). Across synthetic and real single-cell RNA sequencing datasets, WFR-MFM achieves orders-of-magnitude faster inference than a range of existing baselines while maintaining high predictive accuracy, and enables efficient perturbation response prediction on large synthetic datasets with thousands of conditions.

cs.LG↗

arXiv2026

Multi-Task Fine-Tuning Enables Robust Out-of-Distribution Generalization in Atomistic Models

Accurate de novo molecular and materials design requires structure-property models that generalize beyond known regimes. Although pretrained atomistic models achieve strong in-distribution accuracy after fine-tuning, their reliability under out-of-distribution (OOD) conditions remains unclear. We identify a critical failure mode in downstream adaptation: standard fine-tuning induces representation collapse, erasing pretrained chemical and structural priors and severely degrading OOD performance. To address this limitation, we propose multi-task fine-tuning (MFT), which jointly optimizes downstream property prediction with a physically grounded force-field objective inherited from pretraining. This approach preserves essential chemical priors while enabling task-specific adaptation. Across molecular and materials benchmarks, MFT consistently improves OOD generalization, approaching the theoretical limit set by in-distribution accuracy, while outperforming standard fine-tuning, training from scratch, and state-of-the-art task-specific models. These results establish safe adaptation as a central requirement for large atomistic models and position MFT as a practical and data-efficient pathway toward robust molecular and materials discovery.

physics.comp-ph↗