arXiv subjects

Rongjie Lai

Publications and source records attributed to Rongjie Lai.

At least 19 recordsLinked to original sources

arXiv2026

Neural Operator Learning for Collision-Aware Trajectory Planning of Spacecraft Swarms

Satellite constellations require orbital transfers that are both fuel efficient and collision avoidant. Yet, the computational cost of optimization methods traditionally used to plan their trajectories scales poorly with both the number of satellites as well as the number of obstacles to avoid, due to the pairwise safety constraints. In this work, we introduce a permutation-equivariant neural operator for trajectory planning of spacecraft swarms. This neural operator maps distributions of spacecraft initial states, target states, and obstacle initial states to trajectories which avoid collision and conserve fuel. This neural operator output is then paired with a batched Gauss-Newton finish to enforce exact orbital dynamics, and further reduce fuel use. The operator is self-supervised, trained without optimal trajectory labels. When trained on ten spacecraft, the proposed method generalized zero-shot to swarms of 1,000 spacecraft and 11,000 obstacles. The generated trajectories matched a per-agent optimal control solver's accuracy while retaining collision avoidance. Operator learning grounded in physics may offer a fast, scalable alternative to trajectory optimization in the increasingly crowded orbits of the future.

cs.LG↗

arXiv2026

Mode Coverage in Normalizing Flow Boltzmann Generators via Log-Ratio Variation

Normalizing flow Boltzmann generators retain a tractable pushforward density, but training with forward KL depends on target samples that may be biased or omit modes. As a result, a flow can miss target mass while its observed importance weights give a high effective sample size. We introduce the log-ratio variation $\X_ω$, the mean absolute pairwise difference of the target-to-pushforward log-density ratio under a weighting measure $ω$, and use it to define KLXX, a new loss function. Two log-ratio variations are added to the forward KL (denoted by the two X's): one weighted by the target to improve accuracy, the other by a mixture of quench and temper samples with pushforward samples to search candidate modes. We derive the Fisher--Rao gradient flow of KLXX, where both variations contribute nonpositive dissipation, and a fixed-surrogate error bound for KLXX. We use KLXX in an adaptive-staging Boltzmann generator, with importance reweighting at every stage. We bound the sampling error of its inference scheme when the stage weights are essentially bounded, and prove it asymptotically unbiased in the sample size. In the numerical tests, KLXX improves mode coverage over forward KL. It also improves the generator's per-stage diagnostics against the loss that built the schedule. The observables the generator recovers are close to independent references. The log-ratio variations thus supply information that the forward KL loss usually omits.

stat.ML↗

arXiv2026

Self-supervised In-context Operator Learning for Stochastic Mean-Field Control

Stochastic mean-field control (MFC) provides a fundamental framework for coordinating large populations of interacting agents under uncertainty, with a wide range of applications. Existing numerical and deep-learning methods solve one MFC problem instance at a time and must be re-optimized whenever the task changes. In this work, we formulate stochastic MFC as an operator-learning problem and develop, to the best of our knowledge, the first mesh-free, self-supervised neural operator for stochastic MFC. The main challenge is that the diffusion term in the controlled Fokker--Planck equation precludes deterministic transport-map representations. We address this challenge by combining the probability-flow ODE with an invertible normalizing-flow-based transformer, which recasts the dynamics as a deterministic continuity equation and enables closed-form score evaluation through the exact inverse and analytical log-determinant of the normalizing flow, with $\mathcal{O}(d)$ cost per particle for networks of fixed size. Through transformer-based in-context learning, task prompts, represented by compact distribution parameters or raw particle clouds, condition the transport map, enabling a single pretrained operator to solve unseen tasks in one forward pass. The resulting \emph{Normalizing Flow Invertible Solution Transformer} (NFIST) is trained end-to-end by minimizing the stochastic control objective directly, requiring no precomputed numerical solutions for training. We further prove the consistency of the proposed operator-learning formulation with task-by-task optimization. Numerical experiments on stochastic optimal control, Schrödinger bridge, systemic-risk control, and obstacle-avoiding path planning demonstrate effective zero-shot generalization while substantially reducing the computational cost of solving large families of stochastic MFC problems.

math.OC↗

arXiv2026

Training-Free Universal Approximation by Prompting Random Transformers

How expressive is prompting a transformer? Answering this question is important for separating the roles of prompting, architecture, and pretraining in transformer models, and for determining whether task-specific behavior must be stored in model weights or can instead be induced at inference time through the prompt. We show, in an approximation-theoretic sense, that pretraining is optional: a single-layer softmax attention network with random, untrained weights can approximate any Hölder function on a compact manifold when steered by an appropriate soft prompt. Guided by the connection between softmax attention and kernel methods, we construct explicit soft prompts (a prompt per target function, independent of the query) as solutions to linear systems matching attention logits to Gaussian kernel exponents, under which the frozen transformer emulates the classical Nadaraya-Watson kernel estimator. The construction requires only a mild rank condition on the weights, which we show holds almost surely under Gaussian initialization. The prompted network inherits the theoretical guarantees of kernel regression, leading to universal approximation theorems with minimax-optimal rates that depend on the intrinsic dimension. We further quantify the cost of prompting, exposing a tradeoff between the norm of the constructed soft prompt tokens, prompt length, and hidden dimension. Numerical experiments corroborate the constructions and predicted rates.

cs.LG↗

arXiv2026

Scalable Dynamic Optimal Transport via Distributed Linearized ADMM

In this paper, we address two fundamental challenges in the numerical solution of dynamic opti- mal transport (OT) problems. The first challenge arises when the initial and/or terminal densities approach zero and no positive lower bound is available. In this regime, conventional methods may become unstable or computationally inefficient, since the Lipschitz constant of the objective can scale like the reciprocal of the cube of the density. As a result, near-zero regions may lead to slow convergence or even divergence. The second challenge concerns the substantial memory cost of the dynamic formulation, whose discretization over fine spatial and temporal grids requires storing vari- ables across the entire space-time domain. This storage burden quickly becomes prohibitive as the grid is refined or the spatial dimension increases. To overcome the first difficulty, we reformulate the classical discretized dynamic OT problem so that the resulting objective admits an exact proximal mapping. When combined with a linearized alternating direction method of multipliers (LADMM), this reformulation yields an efficient and robust algorithm that remains stable even in challenging settings where the density may vanish. To reduce the memory burden, we further introduce a dis- tributed formulation in which the optimization variables are partitioned across multiple agents. This design substantially lowers the storage requirement for each agent and can also lead to computational acceleration. We validate the proposed framework through numerical experiments in one- and two- dimensional spatial settings under varying levels of difficulty. The results demonstrate the stability, robustness, and scalability of the proposed method in comparison with conventional approaches.

math.OC↗

arXiv2026

Exact Total Variation Minimizers as Non-Oscillatory Limiters in High-Order Methods for Conservation Laws

High-order numerical methods for conservation laws can generate spurious oscillations near discontinuities. We propose to suppress these oscillations by post-processing the numerical solution with a total variation (TV) denoising step that reduces the total variation of the nodal values. For a general analysis operator, the resulting discrete minimization problem need not satisfy the submodularity property that existing exact max-flow algorithms require, so we instead compute the TV minimizer exactly using a differential inclusion algorithm. In one dimension, the differential inclusion algorithm computes the TV minimizer in finitely many steps, requires no tuning of algorithmic parameters, and preserves the total mass. In two dimensions, we apply the 1D algorithm dimension by dimension. We test the method as a post-processing limiter for Fourier pseudospectral methods and a fifth-order finite difference scheme applied to scalar conservation laws, compressible Euler equations, and the two-dimensional incompressible Euler equations.

math.NA↗

arXiv2026

Transformers for Learning on Noisy and Task-Level Manifolds: Approximation and Generalization Insights

Transformers serve as the foundational architecture for large language and video generation models, such as GPT, BERT, SORA and their successors. Empirical studies have demonstrated that real-world data and learning tasks exhibit low-dimensional structures, along with some noise or measurement error. The performance of transformers tends to depend on the intrinsic dimension of the data/tasks, though theoretical understandings remain largely unexplored for transformers. This work establishes a theoretical foundation by analyzing the performance of transformers for regression tasks involving noisy input data near a manifold. Specifically, the input data are in a tubular neighborhood of a manifold, while the ground truth function depends on the projection of the noisy data onto this manifold, referred to as the task-level manifold. We prove approximation and generalization errors which crucially depend on the intrinsic dimension of the task-level manifold. Our results demonstrate that transformers can leverage low-complexity structures in learning task even when the input data are perturbed by high-dimensional noise. Our novel proof technique constructs representations of basic arithmetic operations by transformers, which may hold independent interest.

cs.LG↗

arXiv2026

Understanding In-Context Learning on Structured Manifolds: Bridging Attention to Kernel Methods

While in-context learning (ICL) has achieved remarkable success in natural language and vision domains, its theoretical understanding-particularly in the context of structured geometric data-remains unexplored. This paper initiates a theoretical study of ICL for regression of Hölder functions on manifolds. We establish a novel connection between the attention mechanism and classical kernel methods, demonstrating that transformers effectively perform kernel-based prediction at a new query through its interaction with the prompt. This connection is validated by numerical experiments, revealing that the learned query-prompt scores for Hölder functions are highly correlated with the Gaussian kernel. Building on this insight, we derive generalization error bounds in terms of the prompt length and the number of training tasks. When a sufficient number of training tasks are observed, transformers give rise to the minimax regression rate of Hölder functions on manifolds, which scales exponentially with respect to the prompt length with the exponent depending on the intrinsic dimension of the manifold, rather than the ambient space dimension. Our result also characterizes how the generalization error scales with the number of training tasks, shedding light on the complexity of transformers as in-context kernel algorithm learners. Our findings provide foundational insights into the role of geometry in ICL and novels tools to study ICL of nonlinear models.

cs.LG↗

arXiv2026

Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer

Pre-trained transformers are able to learn from examples provided as part of the prompt without any weight updates, a remarkable ability known as in-context learning (ICL). Despite its demonstrated efficacy across various domains, the theoretical understanding of ICL is still developing. Whereas most existing theory has focused on linear models, we study ICL in the nonlinear regression setting. Through the interaction mechanism in attention, we explicitly construct transformer networks to realize nonlinear features, such as polynomial or spline bases, which span a wide class of functions. Based on this construction, we establish a framework to analyze end-to-end in-context nonlinear regression with the constructed features. Our theory provides finite-sample generalization error bounds in terms of context length and training set size. We numerically validate the theory on synthetic regression tasks.

cs.LG↗

arXiv2026

Improving Clean Accuracy via a Tangent-Space Perspective on Adversarial Training

Adversarial training has proven effective in improving the robustness of deep neural networks against adversarial attacks. However, this enhanced robustness often comes at the cost of a substantial drop in accuracy on clean data. In this paper, we address this limitation by introducing Tangent Direction Guided Adversarial Training (TART), a novel method that enhances clean accuracy by exploiting the geometry of the data manifold. We argue that adversarial examples with large components in the normal direction can overly distort the decision boundary and degrade clean accuracy. TART addresses this issue by estimating the tangent direction of adversarial examples and adaptively modulating the perturbation bound based on the norm of their tangential component. To the best of our knowledge, TART is the first adversarial defense framework that explicitly incorporates the concept of tangent space and direction into adversarial training. Extensive experiments on both synthetic and benchmark datasets demonstrate that TART consistently improves clean accuracy while maintaining robustness against adversarial attacks.

cs.LG↗

arXiv2026

In-Context Operator Learning on the Space of Probability Measures

We introduce \emph{in-context operator learning on probability measure spaces} for optimal transport (OT). The goal is to learn a single solution operator that maps a pair of distributions to the OT map, using only few-shot samples from each distribution as a prompt and \emph{without} gradient updates at inference. We parameterize the solution operator and develop scaling-law theory in two regimes. In the \emph{nonparametric} setting, when tasks concentrate on a low-intrinsic-dimension manifold of source--target pairs, we establish generalization bounds that quantify how in-context accuracy scales with prompt size, intrinsic task dimension, and model capacity. In the \emph{parametric} setting (e.g., Gaussian families), we give an explicit architecture that recovers the exact OT map in context and provide finite-sample excess-risk bounds. Our numerical experiments on synthetic transports and generative-modeling benchmarks validate the framework.

cs.LG↗

arXiv2026

Learn to Evolve: Self-supervised Neural JKO Operator for Wasserstein Gradient Flow

The Jordan-Kinderlehrer-Otto (JKO) scheme provides a stable variational framework for computing Wasserstein gradient flows, but its practical use is often limited by the high computational cost of repeatedly solving the JKO subproblems. We propose a self-supervised approach for learning a JKO solution operator without requiring numerical solutions of any JKO trajectories. The learned operator maps an input density directly to the minimizer of the corresponding JKO subproblem, and can be iteratively applied to efficiently generate the gradient-flow evolution. A key challenge is that only a number of initial densities are typically available for training. To address this, we introduce a Learn-to-Evolve algorithm that jointly learns the JKO operator and its induced trajectories by alternating between trajectory generation and operator updates. As training progresses, the generated data increasingly approximates true JKO trajectories. Meanwhile, this Learn-to-Evolve strategy serves as a natural form of data augmentation, significantly enhancing the generalization ability of the learned operator. Numerical experiments demonstrate the accuracy, stability, and robustness of the proposed method across various choices of energies and initial conditions.

cs.LG↗

arXiv2025

Self-Supervised Amortized Neural Operators for Optimal Control: Scaling Laws and Applications

Optimal control provides a principled framework for transforming dynamical system models into intelligent decision-making, yet classical computational approaches are often too expensive for real-time deployment in dynamic or uncertain environments. In this work, we propose a method based on self-supervised neural operators for open-loop optimal control problems. It offers a new paradigm by directly approximating the mapping from system conditions to optimal control strategies, enabling instantaneous inference across diverse scenarios once trained. We further extend this framework to more complex settings, including dynamic or partially observed environments, by integrating the learned solution operator with Model Predictive Control (MPC). This yields a solution-operator learning method for closed-loop control, in which the learned operator supplies rapid predictions that replace the potentially time-consuming optimization step in conventional MPC. This acceleration comes with a quantifiable price to pay. Theoretically, we derive scaling laws that relate generalization error and sample/model complexity to the intrinsic dimension of the problem and the regularity of the optimal control function. Numerically, case studies show efficient, accurate real-time performance in low-intrinsic-dimension regimes, while accuracy degrades as problem complexity increases. Together, these results provide a balanced perspective: neural operators are a powerful novel tool for high-performance control when hidden low-dimensional structure can be exploited, yet they remain fundamentally constrained by the intrinsic dimensional complexity in more challenging settings.

math.OC↗

arXiv2025

Data-Driven Model Reduction using WeldNet: Windowed Encoders for Learning Dynamics

Many problems in science and engineering involve time-dependent, high dimensional datasets arising from complex physical processes, which are costly to simulate. In this work, we propose WeldNet: Windowed Encoders for Learning Dynamics, a data-driven nonlinear model reduction framework to build a low-dimensional surrogate model for complex evolution systems. Given time-dependent training data, we split the time domain into multiple overlapping windows, within which nonlinear dimension reduction is performed by auto-encoders to capture latent codes. Once a low-dimensional representation of the data is learned, a propagator network is trained to capture the evolution of the latent codes in each window, and a transcoder is trained to connect the latent codes between adjacent windows. The proposed windowed decomposition significantly simplifies propagator training by breaking long-horizon dynamics into multiple short, manageable segments, while the transcoders ensure consistency across windows. In addition to the algorithmic framework, we develop a mathematical theory establishing the representation power of WeldNet under the manifold hypothesis, justifying the success of nonlinear model reduction via deep autoencoder-based architectures. Our numerical experiments on various differential equations indicate that WeldNet can capture nonlinear latent structures and their underlying dynamics, outperforming both traditional projection-based approaches and recently developed nonlinear model reduction methods.

stat.ML↗

arXiv2025

Joint Inference of Trajectory and Obstacle in Mean-Field Games via Bilevel Optimization

Mean field game (MFG) is an expressive modeling framework for systems with a continuum of interacting agents. While many approaches exist for solving the forward MFG, few have studied its \textit{inverse} problem. In this work, we seek to recover optimal agent trajectories and the unseen spatial obstacle given partial observation on the former. To this end, we use a special type of generative models, normalizing flow, to represent the trajectories and propose a novel formulation of inverse MFG as a bilevel optimization (BLO) problem. We demonstrate the effectiveness of our approach across various MFG scenarios, including those involving multi-modal and disjoint obstacles, highlighting its robustness with respect to obstacle complexity and dimensionality. Alternatively, our formulation can be interpreted as regularizing maximum likelihood trajectory learning with MFG assumptions, which improves generalization performance especially with scarce training data. Impressively, our method also recovers the hidden obstacle with high fidelity in this low-data regime.

math.OC↗

arXiv2024

Semi-Supervised Manifold Learning with Complexity Decoupled Chart Autoencoders

Autoencoding is a popular method in representation learning. Conventional autoencoders employ symmetric encoding-decoding procedures and a simple Euclidean latent space to detect hidden low-dimensional structures in an unsupervised way. Some modern approaches to novel data generation such as generative adversarial networks askew this symmetry, but still employ a pair of massive networks--one to generate the image and another to judge the images quality based on priors learned from a training set. This work introduces a chart autoencoder with an asymmetric encoding-decoding process that can incorporate additional semi-supervised information such as class labels. Besides enhancing the capability for handling data with complicated topological and geometric structures, the proposed model can successfully differentiate nearby but disjoint manifolds and intersecting manifolds with only a small amount of supervision. Moreover, this model only requires a low-complexity encoding operation, such as a locally defined linear projection. We discuss the approximation power of such networks and derive a bound that essentially depends on the intrinsic dimension of the data manifold rather than the dimension of ambient space. Next we incorporate bounds for the sampling rate of training data need to faithfully represent a given data manifold. We present numerical experiments that verify that the proposed model can effectively manage data with multi-class nearby but disjoint manifolds of different classes, overlapping manifolds, and manifolds with non-trivial topology. Finally, we conclude with some experiments on computer vision and molecular dynamics problems which showcase the efficacy of our methods on real-world data.

cs.LG↗

arXiv2024

Unsupervised Solution Operator Learning for Mean-Field Games via Sampling-Invariant Parametrizations

Recent advances in deep learning has witnessed many innovative frameworks that solve high dimensional mean-field games (MFG) accurately and efficiently. These methods, however, are restricted to solving single-instance MFG and demands extensive computational time per instance, limiting practicality. To overcome this, we develop a novel framework to learn the MFG solution operator. Our model takes a MFG instances as input and output their solutions with one forward pass. To ensure the proposed parametrization is well-suited for operator learning, we introduce and prove the notion of sampling invariance for our model, establishing its convergence to a continuous operator in the sampling limit. Our method features two key advantages. First, it is discretization-free, making it particularly suitable for learning operators of high-dimensional MFGs. Secondly, it can be trained without the need for access to supervised labels, significantly reducing the computational overhead associated with creating training datasets in existing operator learning methods. We test our framework on synthetic and realistic datasets with varying complexity and dimensionality to substantiate its robustness.

cs.LG↗

arXiv2024

Emulating Complex Synapses Using Interlinked Proton Conductors

In terms of energy efficiency and computational speed, neuromorphic electronics based on non-volatile memory devices is expected to be one of most promising hardware candidates for future artificial intelligence (AI). However, catastrophic forgetting, networks rapidly overwriting previously learned weights when learning new tasks, remains as a pivotal hurdle in either digital or analog AI chips for unleashing the true power of brain-like computing. To address catastrophic forgetting in the context of online memory storage, a complex synapse model (the Benna-Fusi model) has been proposed recently[1], whose synaptic weight and internal variables evolve following a diffusion dynamics. In this work, by designing a proton transistor with a series of charge-diffusion-controlled storage components, we have experimentally realized the Benna-Fusi artificial complex synapse. The memory consolidation from coupled storage components is revealed by both numerical simulations and experimental observations. Different memory timescales for the complex synapse are engineered by the diffusion length of charge carriers, the capacity and number of coupled storage components. The advantage of the demonstrated complex synapse in both memory capacity and memory consolidation is revealed by neural network simulations of face familiarity detection. Our experimental realization of the complex synapse suggests a promising approach to enhance memory capacity and to enable continual learning.

cs.NE↗