Machine Learning
Towards Generalization-Oriented Models for Vehicle Routing Problems with Mixture-of-Experts
Changhao Miao, Yuntian Zhang, Tongyu Wu et al. (+2 more)
In recent years, Deep Reinforcement Learning (DRL) has achieved substantial progress on Vehicle Routing Problems (VRPs). However, existing DRL-based methods are typically trained on instances generated from a uniform distribution, which limits their performance under real-world distribution shifts. In this paper, we aim to develop a generalization-oriented model that partitions the policy network into multiple modules and adaptively recombines modules to form specific policies during inference. Specifically, we propose Residual Refined Experts with Instance-level Gating (R2E-IG) to improve cross-distribution generalization. Our contributions are threefold: (1) We introduce a Residual Refined Expert (R2E) architecture that enhance expert expressiveness via residual refinement; (2) We design an instance-level gating mechanism that learns distribution-aware instance representations and routes inputs to suitable modules; (3) We propose a mixed-distribution training mechanism equipped with Dynamic Weight Adaption (DWA), which dynamically reweights training data from different distributions to emphasize more informative ones. Extensive experiments show that R2E-IG achieves competitive performance against state-of-the-art baselines on both in-distribution and out-of-distribution instances across synthetic and
benchmark datasets. Moreover, R2E-IG is generic and can be easily integrated into existing DRL-based methods to further improve performance.
Read the full abstract ▾
All authors Changhao Miao, Yuntian Zhang, Tongyu Wu, Fang Deng, Chen Chen
Machine Learning · Cross-listed from cs.AI
Beyond a Single Direction: Chain-of-Thought Disrupts Simple Steering of Refusal
Kia-Jüng Yang, Dominik Meier, Jiachen Zhao et al. (+2 more)
Large
reasoning models (LRMs) generate chain-of-thought (CoT) traces before producing final outputs, introducing a dynamic internal state that may complicate control mechanisms such as refusal. Unlike instruction-tuned LLMs, where refusal is mediated by a single directional subspace, refusal in large
reasoning models (LRMs) additionally depends on the CoT. In DeepSeek-R1-Distill-LLaMA-8B, activation steering reverses refusal in only 39% of cases when the CoT is kept fixed, but removing the CoT entirely increases this to 70%, indicating that the CoT actively reinforces refusal. In a two-stage intervention where the model regenerates its CoT under activation steering, refusal is reversed in 94% of cases, while the resulting CoT alone retains 48% of this effect even after steering is removed. This suggests that the CoT can carry and reconstruct the compliance signal independently. These findings indicate that refusal in LRMs is jointly encoded in residual stream activations and CoT. This joint activation makes LRM more robust against activation-level interventions alone, but exposes CoT to a possible alternative surface attack.
Read the full abstract ▾
All authors Kia-Jüng Yang, Dominik Meier, Jiachen Zhao, Terry Ruas, Bela Gipp
Machine Learning
Adversarial Training for Robust Coverage Network under Worst-case Facility Losses
Changhao Miao, Yuntian Zhang, Tongyu Wu et al. (+2 more)
The Maximal Covering Location-Interdiction Problem (MCLIP) is a classic bi-level optimization problem, which is fundamental to resilient infrastructure planning yet remains computationally intractable. Specifically, the upper level determines facility locations to maximize coverage, while the lower level executes worst-case interdiction to minimize the coverage. The strong coupling between the upper and lower levels, combined with their respective high combinatorial complexity, renders traditional methods ineffective. To bridge this gap, we propose a Dual-
Agent Deep Reinforcement Learning (DADRL) framework based on
adversarial learning, comprising a location
agent corresponding to the upper level and an interdiction
agent corresponding to the lower level. Our contributions are threefold: (1) The location
agent is trained simultaneously against an evolving interdiction
agent, making it effectively capture the dynamic competitive interplay between the upper and lower levels; (2) To fully exploit the learned capabilities of the interdiction
agent, we propose a Surrogate-based Ensemble Inference Strategy that utilizes the trained interdiction
agent as a high-fidelity surrogate to guide the decisions of location
agent; (3) Extensive experiments on synthetic and real-world
datasets demonstrate that our approach achieves superior computational efficiency while maintaining highly competitive solution quality compared to other baselines. Furthermore, our DADRL framework is model-agnostic to network structures, while its underlying
adversarial learning paradigm demonstrates strong potential for solving other bi-level optimization problems.
Read the full abstract ▾
All authors Changhao Miao, Yuntian Zhang, Tongyu Wu, Fang Deng, Chen Chen
Machine Learning
Time Series Causal Discovery via Context-Conditioned and Causality-Augmented Pretraining
Biao Ouyang, Tengxue Zhang, Zhihao Zhuang et al. (+3 more)
Causal discovery from time series is critical for many real-world applications, such as tracing the root causes of anomalies. Existing approaches typically rely on
dataset-specific optimization, making it difficult to transfer their causal discovery capabilities to new time series governed by diverse causal mechanisms. In this paper, we propose PTCD, a novel Pretraining framework for Time-series Causal Discovery, which improves cross-task generalization through context-conditioned modeling and transferable causal augmentation. To model complex temporal causal dependencies, PTCD employs a dual-scale iterative
attention mechanism to capture window-level causal relationships, and a Gaussian mixture with a context-level routing mechanism to handle heterogeneous exogenous distributions. To further address distribution shifts across causal graphs, PTCD adopts a pretraining paradigm on synthetic
datasets that integrates intervention-based learning and a causal mixup strategy, promoting stable causal discovery and stronger generalization. Extensive experiments on multiple real-world out-of-distribution (OOD)
datasets demonstrate that PTCD excels in both causal discovery and root cause identification.
Read the full abstract ▾
All authors Biao Ouyang, Tengxue Zhang, Zhihao Zhuang, Yang Shu, Chenjuan Guo, Bin Yang
Machine Learning
Localizing Memorized Regions in Diffusion Models via Coordinate-Wise Curvature Differences
Gwangho Kim and Sungyoon Lee
Diffusion models can unintentionally memorize training samples, raising concerns about
privacy and copyright. While recent methods can detect memorization, they often rely on global or model-specific signals and provide limited insight into where memorization appears within a generated image. We provide a
geometric characterization of local memorization as a coordinate-wise variance collapse. However, such collapse can also arise from intrinsic data constraints rather than overfitting. To isolate overfitting-driven memorization, we propose
curvature-difference methods that subtract the
curvature of an underfitted baseline, either the unconditional model or a less-trained version of itself. We further derive a score-difference proxy that provides a
geometric explanation for the widely used score-difference-based detection metric. Experiments on Stable
Diffusion, evaluated against ground-truth memorization masks, show that our method outperforms the prior
attention-based localization method. Code is available at https://github.com/Gwangho99/mem-curv-diff.
Read the full abstract ▾
Machine Learning
Stabilizing Recurrent Dynamics for Test-Time Scalable Latent Reasoning in Looped Language Models
Xiao-Wen Yang, Ziyu Han, Xi-Hua Zhang et al. (+4 more)
Looped
Language Models (LoopLMs) enable efficient latent
reasoning through depth recurrence, yet exhibit unreliable test-time scaling behavior: performance often peaks at a certain iteration depth and then collapses with further recurrence. Through latent dynamics analysis, we find an inherent trade-off between stability and effectiveness in existing architectures and strategies. By conceptualizing
reasoning as uncertainty reduction, we propose that convergence toward stable fixed points while preserving effectiveness represents a promising way. To this end, we propose
STARS (STAbility-driven Recurrent Scaling), a training framework that constrains latent states to approach asymptotically stable fixed points. This is realized via efficient Jacobian Spectral Radius Regularization with random loop sampling, enabling
STARS to maximize effectiveness while ensuring rigorous stability. Experiments on arithmetic tasks show that
STARS achieves reliable test-time scaling, and on complex mathematical
reasoning it substantially mitigates performance degradation as recurrence depth increases while also improving peak performance.
Read the full abstract ▾
All authors Xiao-Wen Yang, Ziyu Han, Xi-Hua Zhang, Wen-Da Wei, Jie-Jing Shao, Lan-Zhe Guo, Yu-Feng Li
Machine Learning
APEX: Amplitude Anchors and Phase Priors for Target-Scarce Higher-Frequency Wave Prediction
Yifan Sun, Lei Cheng, Sijie Chen et al. (+3 more)
Learning-based surrogates have become increasingly effective for wave-field prediction, and neural operators in particular have shown strong performance within observed frequency regimes. However, higher-frequency prediction under scarce target supervision remains comparatively underexplored, especially in wave problems where higher-frequency data are substantially more expensive to simulate or measure than lower-frequency data. A central difficulty is that cross-frequency transfer is inherently asymmetric: coarse amplitude structure remains relatively stable across frequencies, whereas phase-sensitive oscillatory structure deteriorates much more rapidly as frequency increases. Motivated by this asymmetry, we propose APEX, Amplitude-anchored and Phase-prior-guided Enhancement from eXtrapolated coarse predictions, a framework for target-scarce higher-frequency wave-field prediction. A lower-frequency neural operator first provides a coarse prediction in the target-frequency regime, from which we retain only the amplitude as a transferable structural anchor. A conditional flow-matching enhancer then reconstructs the target higher-frequency field under the guidance of a Green's-function-inspired phase prior. Experiments on SimpleWave, Helmholtz, and Maxwell
benchmarks show that APEX consistently outperforms direct lower-to-higher extrapolation, target-adapted operator, and joint generative baselines under limited target-frequency supervision. Our results suggest that reliable higher-frequency prediction of oscillatory wave fields should not rely on direct end-to-end transfer of the full complex field, but instead on explicitly reusing transferable coarse structure while separately recovering the missing oscillatory detail.
Read the full abstract ▾
All authors Yifan Sun, Lei Cheng, Sijie Chen, Ting Zhang, Jianlong Li, Shikai Fang
Machine Learning
MTL-FNO: A Lightweight Multi-Task Fourier Neural Operator for Sparse Field Reconstruction
Siyu Ye, Shihang Li, Zhiqiang Gong et al. (+4 more)
Efficient onboard multi-field sparse reconstruction is essential for the autonomous operation of aerospace vehicles. While existing deep learning models exhibit promise for single-field reconstruction, deploying multiple independent models leads to prohibitive model size growth and fails to exploit cross-field correlations, particularly under few-shot conditions. To address these challenges, we first propose a lightweight multi-task
Fourier neural operator (MTL-FNO), an end-to-end joint training framework based on hard parameter sharing. In each layer, the parameters are divided into shared and task-specific components to capture common features across fields while preserving task-specific characteristics. Moreover, the task-specific fine-tuning parameters are implemented as low-rank terms, achieving substantial model compression. Second, to address the difficulty of co-optimizing shared and task-specific parameters along with their real and imaginary parts, we revisit the FNO's spectral weight from a polar-form perspective and devise a physically meaningful decoupled optimization scheme. Specifically, we apply polar decomposition to slice-wise disentangle the spectral weight into a
unitary tensor encoding phase information and a positive semi-definite tensor characterizing amplitude. By decoupling the optimization of phase and amplitude, our method can effectively mitigate tasks conflict. Meanwhile, to preserve
unitary geometric fidelity during training, the Cayley transform is introduced to reparameterize the
unitary tensor, converting the constrained optimization problem to an unconstrained one. Finally, the effectiveness of the proposed method under few-shot conditions is validated on two representative engineering cases. Results show that MTL-FNO achieves accuracy comparable to or even surpassing that of standard FNO, while reducing total model size by 76% and 60%, respectively.
Read the full abstract ▾
All authors Siyu Ye, Shihang Li, Zhiqiang Gong, Benrong Zhang, Weien Zhou, Yiyong Huang, Wen Yao
Machine Learning
Image Feature Fusion-based Federated Client Unlearning (FCU)
Hangyi Shen, Yizhi Pan, Tiansuo Li et al. (+2 more)
Major data protection regulations all mention the "right to be forgotten," and that's what pushed
federated unlearning (FU) techniques forward. But one stubborn issue remains: catastrophic forgetting–you erase the target knowledge, yet somehow you also end up throwing out essential retained knowledge, which then hurts the model's global generalization. To get a better balance between unlearning effectiveness and generalization ability, we propose something called Image Feature Fusion-based
Federated Client Unlearning (IFF-FCU). The idea is to bring in a linear Image Feature Fusion mechanism (Mixup) that dynamically creates mixed samples, bridging the gap between forget-distribution and retain-distribution. What this strategy does isn't just deleting a few discrete data points–it theoretically widens and regularizes the forgetting boundary. We ran extensive experiments on medical imaging
benchmarks (RSNA-ICH and ISIC2018), and the results show that our approach achieves reasonably good unlearning. For instance, on the ICH
dataset, IFF-FCU achieves a highly competitive Error deviation from the retrained gold standard, demonstrating robust improvements over existing baselines.
Read the full abstract ▾
All authors Hangyi Shen, Yizhi Pan, Tiansuo Li, Weiqi Jiang, Guanqun Sun
Machine Learning · Cross-listed from stat.ML
Transformers Can Learn Posterior Predictive Distributions In-Context
Gyeonghun Kang, Changwoo J. Lee and Xiang Cheng
Prior-data fitted networks (PFNs) have recently emerged as a powerful approach for
Bayesian prediction tasks, approximating the posterior predictive distribution (PPD) through in-context learning. Despite their strong empirical performance and ability to go beyond point predictions, theoretical understandings of the algorithmic capability of
transformers to learn distributions in context are still lacking. Focusing on Gaussian process regression problems, we show by construction that
transformers can implement a gradient descent
algorithm targeting the posterior predictive mean and variance, followed by nonlinear mappings that yield binned probabilities of PPD. We study the error bounds of the approximated PPD in terms of
attention depth and bin resolution. Based on these results, we further demonstrate the key role of normalization and the choice of
attention depth in enabling the extrapolation abilities of
transformers beyond the pretraining sample size range. We conduct simulations that corroborate our findings, providing insight into the expressivity of PFNs targeting PPDs and how architectural choices may influence generalization capabilities.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
The Need for an External Observer Formalizing the Sufficiency Gap: A Mathematical Extension of Mixture Identifiability and Contextual Grounding in Sequence Models
Francesco Corielli
We construct a binary mixed-regime process with one deterministic textual regime and one random regime governed by an unobserved latent state. Even an ideal infinite-capacity sequence predictor that exactly recovers the text-only marginal law can become overconfident when the observed prefix is compatible with the wrong latent regime. The resulting
entropy difference is not an ordinary optimization error; it is a sufficiency gap caused by marginalization over an unobserved state. We then formalize
retrieval, tool use, and external grounding through an auxiliary binary signal with fidelity $\gamma \in [1/2,1]$. The resulting
Bayesian update yields a contextual dominance threshold: a corrective signal reverses the posterior odds induced by the textual history exactly when its fidelity exceeds the text-only posterior weight assigned to the misleading regime. This threshold reduces, but does not generally eliminate, the sufficiency gap; complete closure requires perfect revelation of the relevant latent state or an equivalent verification mechanism. The analysis clarifies why temperature scaling cannot restore missing context, why grounding mechanisms must be both informative and learnably usable by the model, and why autonomous sequence models require structurally decoupled observers or verifiers in high-stakes domains.
Read the full abstract ▾
Machine Learning
SL-BiLEM: Structured Learnable Behavior-in-the-Loop Epidemic Modeling for Forecasting and Policy Evaluation
Haochun Wang, Sendong Zhao, Jingbo Wang et al. (+3 more)
Epidemic
forecasting faces a fundamental challenge: human behavior dynamically responds to disease spread, creating feedback loops that induce distribution shifts at policy intervention points. This renders data-driven models unreliable under distribution shift. We propose SL-BiLEM (Structured Learnable Behavior-in-the-Loop Epidemic Model), leveraging physical constraints as regularization for robust extrapolation. The framework decomposes effective transmission as $\beta_{\text{eff}}(t,g) = \beta_0(g) \times m_{\text{policy}}(t) \times m_{\text{media}}(t) \times m_{\text{comp}}(t,g)$, where monotonicity, smoothness, and bounded-jump constraints on the learned compliance function maintain predictive validity under novel policy regimes. Beyond
forecasting, SL-BiLEM enables counterfactual analysis for intervention decision support. We validate
forecasting on three real-world
datasets (cruise ship, school influenza, and school-district COVID-19 surveillance) and evaluate counterfactual recovery on synthetic
benchmarks with known ground truth. SL-BiLEM demonstrates: (1) 76% improvement over neural-mechanistic baselines, with only 53% OOD degradation versus 1142% for neural baselines under policy-induced shift; (2) 100% bootstrap CI coverage across 27 synthetic counterfactual experiments; and (3) Treatment Effect Accuracy exceeding 0.85. These results establish SL-BiLEM as an interpretable tool for public health decision-makers seeking accurate prediction and principled intervention planning.
Read the full abstract ▾
All authors Haochun Wang, Sendong Zhao, Jingbo Wang, Yanrui Du, Ting Liu, Bing Qin
Machine Learning · Cross-listed from econ.TH
Proper Calibeating
Dean P. Foster and Sergiu Hart
The classic concept of "calibrated forecasts" and its more recent refinement, "calibeating," are defined with respect to the standard quadratic scoring rule. We extend these notions to the class of proper scoring rules (for which the true distribution is an optimal forecast) and define proper calibration and proper calibeating by requiring the corresponding guarantees to hold uniformly over all bounded proper scoring rules. We first establish that calibration always implies proper calibration, whereas calibeating need not imply proper calibeating. Second, we show how to guarantee proper calibeating and proper multicalibeating; in particular, complete calibeating—a strong form of calibeating that calibeats the joint binning—is always proper. Finally, we consider decision-making under uncertainty, where one best replies to the forecasts. We establish that (proper) calibration is equivalent to universal no regret, and (proper) complete calibeating is equivalent to universal no regret together with universal subsuming of the reference sequence, where universal refers to all bounded utility functions.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CV
Rotation-Invariant Spherical Watermarking via Third-Order SO(3) Representation Coupling
Pengzhen Chen, Yanwei Liu, Xiaoyan Gu et al. (+3 more)
Reliable watermarking of panoramic imagery is fundamentally challenged by arbitrary 3D rotations. As panoramas are defined on the sphere, they naturally transform under the action of $SO(3)$, rendering conventional planar representations and augmentation-based robustness strategies inadequate and devoid of theoretical guarantees. To address this, we formulate panoramas as spherical signals and leverage $SO(3)$ representation theory to derive provably rotation-invariant descriptors. While spherical harmonic coefficients transform equivariantly under rotations, the natural invariant constructions are typically limited to zeroth-order statistics which eliminate directional information and severely constrain
embedding capacity. In this work, we introduce a principled third-order invariant construction by coupling higher-order $SO(3)$ irreducible representations via tensor products and projecting onto the trivial representation. This yields a spherical invariant bispectrum that preserves phase information while remaining strictly rotation-invariant. Leveraging this property, we embed watermarks into higher-order spherical harmonic coefficients and recover them from invariant bispectral scalars, enabling reliable extraction under arbitrary 3D rotations. We provide a theoretical proof of $SO(3)$ invariance for it and demonstrate experimentally its near-perfect robustness to continuous rotations while maintaining high visual fidelity.
Read the full abstract ▾
All authors Pengzhen Chen, Yanwei Liu, Xiaoyan Gu, Antonios Argyriou, Wu Liu, Weiping Wang
Machine Learning
Model Merging on Loss Landscape: A Geometry Perspective
Juanwu Lu, Anand Bhaskar, Brian Axelrod et al. (+2 more)
Model merging offers a promising avenue for knowledge integration and parallel development without retraining. Yet, existing methods either ignore the geometry of the loss landscape or rely on intractable full-space Hessian approximations. We propose EpiMer, a framework that casts model merging as solving the Fréchet mean on a Riemannian
manifold and restricts the computation to a low-rank subspace spanned by the task vectors. With the expected Hessian as the metric, we reveal a connection between local
curvature and epistemic uncertainty of the parameters. Our theoretical analysis decomposes the merging error bound into the subspace Fréchet variance and the residual energy, and provides a closed-form characterization of when
curvature-aware merging provably outperforms flat-geometry methods. In addition, our framework unifies both
curvature-aware methods and recent spectral methods as special cases of the subspace Fréchet mean with different
geometric metrics. Merging fine-tuned CLIP-ViT models on eight image classification tasks, Epistemic Merging strictly outperforms the baselines on all three CLIP-ViT backbones at matched rank, improving the across-task average accuracy and worst-task accuracy on every backbone.
Read the full abstract ▾
All authors Juanwu Lu, Anand Bhaskar, Brian Axelrod, Ekaterina Tolstaya, Tristan Emrich
Machine Learning
Self-Improvement Imitation with Biologically Guided Search for Protein Design Under Oracle Budgets
Ashima Khanna and Dominik Grimm
Protein sequence optimization under tight oracle budgets requires methods that explore vast combinatorial spaces while making each evaluation informative. Existing reinforcement learning and off-policy generative approaches often degrade under surrogate noise, and position-agnostic mutation proposals risk disrupting functionally critical residues. We introduce SILO, a trajectory-level self-improvement imitation framework for oracle-budgeted
protein design. SILO uses a hierarchical edit policy that decomposes each mutation into a position choice followed by a residue choice. In each active-learning round, the policy samples candidate trajectories via incremental
stochastic beam search without replacement (SBS), and a UCB-based proxy ensemble, combined with an alanine-scan fitness score (AFS), selects candidates with functionally relevant edits for in silico oracle evaluation. The policy is then updated by next-action cross-
entropy imitation on the round's best oracle-labeled trajectories, avoiding value-function estimation. Across eight reproduced
protein fitness landscapes and five strong baselines from prior work, SILO achieves the highest maximum and top-100 mean fitness on 8 of 8 landscapes within our evaluations, often exhibiting faster early-stage improvement. In low-data and noisy-proxy stress tests on two landscapes per setting, SILO remains competitive or best when several baselines degrade. Ablations show that SBS with AFS account for much of the gains, with iterative imitation providing additional improvement. Code is available at: https://github.com/grimmlab/SILO.git
Read the full abstract ▾
Machine Learning
Beyond Trajectory-Level Attribution: Graph-Based Credit Assignment for Agentic Reinforcement Learning
Xin Cheng, Shuo He, Lang Feng et al. (+4 more)
Group-based reinforcement learning (RL) methods have achieved remarkable success in improving the performance of large
language models (LLMs) and have been rapidly extended to agentic tasks. However, their credit assignment relies heavily on coarse-grained trajectory-level attribution according to final outcomes, making it difficult to capture the contribution of individual steps, such as valuable steps obscured within failed trajectories. To uncover latent information and enable more faithful step-level credit assignment, we propose Graph-based Group Policy Optimization (GraphGPO), which first aggregates all rollout trajectories into a unified state-transition graph and then estimates the distance from each state to the task goal using the global information encoded in the graph. Finally, GraphGPO assigns credit to each edge by estimating a graph-based advantage, based on how much the transition reduces the distance to the task goal. In this way, GraphGPO significantly improves training efficiency and achieves state-of-the-art performance across a range of challenging
benchmarks.
Read the full abstract ▾
All authors Xin Cheng, Shuo He, Lang Feng, HaiYang Xu, Ming Yan, Lei Feng, Bo An
Machine Learning · Cross-listed from stat.ML
CART Random Forests as Sequential Allocation over Random Opportunity Sets: A Stochastic-Control Theory of Ensemble Risk
Tianxing Mei, Yingying Fan, Mingming Leng et al. (+1 more)
CART random forests are among the most widely used modern predictive methods, with well-documented empirical success. Yet, at the mechanistic level, the
algorithm is often treated as a black box because of its complexity. In this paper, we develop a
stochastic-control perspective on feature-subsampled CART random forests, named CART random opportunity-set allocation (CART-ROSA). At each node, the random subset of features is interpreted as a random feasible action set, and the CART split rule as a masked-action allocation policy. This policy induces a controlled
stochastic process over informative split-count states, whose terminal law determines both single-tree error and cross-tree interaction terms in the forest mean squared error (MSE). Such representation opens the black box of CART-forests by separating two design levers: the informative-opportunity rate induced by feature subsampling, and the contraction strength from the within-mask split policy. We establish that the CART policy is locally stabilizing: it contracts imbalances in informative split allocations and concentrates terminal tree geometry. At the system level, however, it can be globally suboptimal for the forest objective. Specializing to the linear model, we derive the MSE risk expansion explicitly. Our results show how an operations-research perspective makes tractable a theoretical gap difficult to access from the standard algorithmic description of CART forests.
Read the full abstract ▾
All authors Tianxing Mei, Yingying Fan, Mingming Leng, Jinchi Lv
Machine Learning · Cross-listed from cs.AI
MemFail: Stress-Testing Failure Modes of LLM Memory Systems
Ishir Garg, Neel Kolhe, Dawn Song et al. (+1 more)
Large
language model (LLM)
agents increasingly rely on external
memory systems to remain consistent across long-horizon interactions, but little empirical work has been done to understand the specific failure modes and design choices that these systems present. Existing
benchmarks report aggregate question-answering accuracy and treat
memory systems as black boxes, making it impossible to attribute an incorrect answer to a particular failure mode of the system. We introduce MemFail, a diagnostic
benchmark that isolates the failure modes of modern LLM
memory systems. We begin by formalizing
memory systems as the composition of three canonical operations – summarization, storage, and
retrieval – and identify the potential failure modes induced by each. Based on these hypothesized failure modes, we construct five
datasets spanning four tasks, each adversarially designed to test a specific operation of a
memory system. Using these
datasets, we evaluate four state-of-the-art
memory systems on MemFail and demonstrate how MemFail can be used to empirically understand the tradeoffs induced by differences in
memory system architectures.
Read the full abstract ▾
All authors Ishir Garg, Neel Kolhe, Dawn Song, Xuandong Zhao
Machine Learning
WINDQuant: Weight-Informed Neural Decision-Making for Global Mixed-Precision LLM Quantization
Phong Nam Huu Nguyen, Khoi M. Le, Cong-Duy T Nguyen et al. (+3 more)
Quantization is an effective approach to reduce the
memory footprint and inference cost of large
language models (LLMs), yet maintaining performance in the ultra-low-bit regime remains challenging. Existing post-training methods often suffer from severe accuracy degradation, while quantization-aware training requires costly retraining and additional resources. Moreover, most mixed-precision strategies rely on coarse-grained or heuristic sensitivity analysis that overlooks fine-grained variations within weight
matrices. We propose WINDQuant, a reinforcement-learning-based allocation controller for ultra-low-bit LLM quantization. Rather than introducing another low-level quantization operator, WINDQuant learns how to assign bit-widths and quantization treatments to fine-grained column chunks under a global storage budget. By operating at the column-chunk level, WINDQuant enables flexible and fine-grained precision assignment within layers under a global target bit-width. The implementation combines PPO with activation-aware calibration, lightweight per-unit quantizer fitting, and explicit effective-bit accounting of the learned mixed-precision plan. Experiments on LLaMA models demonstrate that WINDQuant achieves competitive performance in ultra-low-bit settings while reducing optimization overhead relative to retraining-based approaches, highlighting reinforcement learning as a practical controller for adaptive mixed-precision quantization.
Read the full abstract ▾
All authors Phong Nam Huu Nguyen, Khoi M. Le, Cong-Duy T Nguyen, Anh Tuan Luu, Thong Thanh Nguyen, Tho Quan
Machine Learning · Cross-listed from cs.CL
Why Prompt Optimization Works, and Why It Sometimes Doesn't: A Causal-Inspired Edit-Level Analysis
Shuzhi Gong and Hechuan Wen
Automated
prompt optimization methods (e.g., DSpy, TextGrad) can substantially improve the performance of large
language model (LLM), however, their generalization ability across different tasks remains underperformed. In practice, the superiority of the optimized
prompt on one
benchmark often fails to transfer to another, and this limitation persists even when switching across different LLM backbones. To investigate the underexplored sources of heterogeneity in
prompt performance, we conduct a causal inference-inspired observational analysis of optimized
prompts across a diverse set of optimization frameworks, LLM backbones, and NLP
benchmarks. To achieve the goal, we build upon the propensity-adjusted associational analysis together with multiple complementary representations of
prompt edits, where the consistent task-conditioned edits patterns are identified. We find that complexity-increasing and meta-instructional edits are negatively associated with mathematical and multi-hop
reasoning performance, whereas step-by-step and meta-cognitive edits improve logical and sequential
reasoning tasks. These effects are robust across cognitive-load annotations, surface-level text features, and edit-motif analyses, and can generalize across optimization frameworks. Overall, these results indicate that
prompt optimization failures arise from systematic interactions between edit families and task characteristics rather than random optimization artifacts, providing feature-level characterization of optimizer behavior and motivating future task-conditioned optimizer design.
Read the full abstract ▾
Machine Learning
Bilevel Optimization over Saddle Points of Zero-Sum Markov Games
Zihao Zheng, Irwin King and Songtao Lu
Reinforcement learning (RL) often has a hierarchical structure, where an upper-level (UL) learner selects model parameters and a lower-level (LL) decision-making process responds, naturally leading to a bilevel optimization problem. Most existing bilevel RL methods assume a single-policy LL Markov decision process (MDP), and therefore fail to capture competitive structures arising in applications such as incentive design, where multiple policies interact. We study bilevel optimization problems in which the LL problem is a regularized min-max zero-sum Markov game and the UL objective is optimized through the saddle-point
equilibrium induced by the LL game. In this work, we propose penalty-augmented Nikaido-Isoda descent-ascent (PANDA), a penalty-based first-order policy-gradient method based on the Nikaido-Isoda function. By exploiting the min-max game structure, PANDA avoids computing UL hypergradients and does not require second-order information. We prove that PANDA converges to stationary points without convexity assumptions on either the UL or LL objectives. Moreover, PANDA reaches an $\epsilon$-stationary point in $\tilde{\mathcal{O}}(\epsilon^{-1})$ iterations with sample complexity $\tilde{\mathcal{O}}(\epsilon^{-3})$, matching the best-known rates for bilevel RL with single-policy LL MDPs. Experiments demonstrate the superior performance of PANDA over closely related baselines.
Read the full abstract ▾
Machine Learning
More Expressive Feedforward Layers: Part I. Token-Adaptive Mixing of Activations
Mingze Wang, Jinbo Wang, Yikuan Xia et al. (+2 more)
Feedforward network (FFN) layers account for a large fraction of parameters and nonlinear expressivity in
Transformer-based large
language models (LLMs). Despite the
evolution from ReLU and GELU to gated variants such as SwiGLU, most FFN designs still use a single fixed activation function, applying the same nonlinear transformation to all
tokens. In this work, we propose Mixture of Activations (MoA), a
token-adaptive FFN design that mixes a dictionary of activation functions using lightweight input-dependent gates while sharing the same linear projections. As an input-independent counterpart, we also introduce learnable activations (LA), which form linear combinations of activation functions for both ReLU-type and SwiGLU-type FFNs. Theoretically, we establish strict finite-width expressive separations among fixed-activation FFNs, LA, and MoA: LA strictly contains fixed-activation FFNs, while MoA strictly contains LA, with the additional expressivity arising from input-dependent nonlinear hybridization. Empirically, we evaluate MoA through extensive pre-training experiments on dense and MoE
language models ranging from 0.12B to 2B parameters under different
token budgets, optimizers, and learning rate schedules. MoA consistently achieves lower terminal loss and exhibits more favorable scaling behavior than well-tuned baselines, with minimal parameter and computational overhead. These results suggest that
token-adaptive activation mixing is a simple and effective mechanism for improving FFN expressivity in LLMs.
Read the full abstract ▾
All authors Mingze Wang, Jinbo Wang, Yikuan Xia, Kai Shen, Shu Zhong
Machine Learning · Cross-listed from eess.SY
Sample Complexity of Policy Gradient for Log-Growth Control
Qiuhua Pan, Yukai Shen, Liwei Zhang et al. (+2 more)
We study the sample complexity of policy gradient for log-growth control – the problem of learning, from observed state transitions, a feedback gain that optimally stabilizes a scalar linear system driven through a multiplicative-noise actuation channel. The objective $J(K) = \mathbb{E}[\log|1+BK|]$ is the top Lyapunov exponent of the closed loop. This problem carries a structural difficulty we call the cusp obstruction: the optimal gain $K^*$ always places the noise singularity $b_{\rm sing}(K) = -1/K$ in the interior of the support. At this singular optimum the policy gradient exists only as a Cauchy principal value, not as a Lebesgue integral, and the natural single-sample gradient
estimator has infinite variance. Standard first-order
stochastic-optimization analysis is thus inapplicable at the optimum, and merely smoothing the objective does not resolve the difficulty. The obstruction, however, has an exploitable
symmetry: the Cauchy kernel is an odd function of the displacement from the moving pole, so pairing each observation with its reflection through the pole cancels the divergent part. This one cancellation simultaneously controls the population
curvature, the gradient-
estimator variance, and the bias incurred when the noise density is estimated. Combining these bounds with a closed-form single-transition gradient oracle, we prove that projected mini-batch policy gradient, initialized in any compact subset of the stabilizing region, attains total sample complexity $\tilde{O}(1/\eta)$ when the noise density is known and $\tilde{O}(\eta^{-(2s+1)/(2s)})$ when it must be estimated, for $C^s$ noise densities with $s \geq 2$.
Read the full abstract ▾
All authors Qiuhua Pan, Yukai Shen, Liwei Zhang, Cailian Chen, Xinping Guan
Machine Learning
RT-Lynx: Putting GEMM Sparsity in the Right Place for Diffusion Models
Xing Cong, Hanlin Tang, Kan Liu et al. (+3 more)
Diffusion Transformers (DiT) achieve strong performance in image generation but incur substantial inference costs. While prior work has reduced this cost via quantization and
distillation, semi-structured sparsity, which can nearly halve FLOPs, remains underexplored. A key reason is that most existing approaches focus on weight sparsification, and
pruning 50% of the weights can remove critical model capacity and degrade generation quality. Our study, however, shows that DiT activations are intrinsically sparse and significantly more robust to N:M semi-structured sparsification than weights. Motivated by this observation, we advocate a paradigm shift from weight sparsification to activation sparsification. We propose RT-Lynx, which applies N:M sparsification to activations and incorporates error-compensation techniques to mitigate accuracy loss. We further implement highly optimized CUDA kernels tailored to this setting, achieving up to a 1.55x speedup on average in linear layers. Extensive experiments across multiple
diffusion models demonstrate that our method preserves the generation quality of the original models while substantially accelerating inference.
Read the full abstract ▾
All authors Xing Cong, Hanlin Tang, Kan Liu, Tao Lan, Lin Qu, Chenhao Xie
Machine Learning · Cross-listed from stat.AP
Data-driven sparse identification of governing PDEs via knockoff filters and multi-criteria trade-offs
Pongpisit Thanasutives, Naichang Ke and Yoshinobu Kawahara
We propose KO-PDE-IDENT, a data-driven framework for identifying parsimonious partial differential equations (PDEs) with false discovery rate (FDR) control. PDE discovery from noisy observations is often hindered by extreme multicollinearity among candidate terms, which causes typical sparse-regression methods to select spurious terms. To address this problem, KO-PDE-IDENT initially mines a support set of potential candidate terms via model-X knockoff filters with finite-sample FDR control, then refines and ranks the surviving PDE alternatives. The framework integrates three components. First, knockoff feature statistics are constructed by coupling $\ell_{0}$-constrained adaptive best-subset selection with SHapley Additive exPlanations (SHAP), yielding an effective and computationally efficient difference statistic. Second, a recursive feature elimination (RFE) procedure removes terms whose marginal contributions are dispensable and assesses statistical necessity through knockoff-perturbed hypothesis testing. Third, the final model selection is formulated as a multi-criteria decision-making (MCDM) problem, where the optimal governing equation is the alternative that best balances a wide range of criteria such as predictive accuracy, model complexity and coefficient uncertainty. We evaluate KO-PDE-IDENT on five canonical PDEs under severe noise corruption. Empirical results show that our framework can exactly recover the true PDE structure, eliminating false discoveries while retaining all true underlying terms, with low coefficient estimation error.
Read the full abstract ▾
Machine Learning
PIDM-DP: Physics-Informed Diffusion with Dormand-Prince Integration for Chaotic System Identification and State Reconstruction across Multiple Dynamical Regimes
Shailendra Dabral
Reconstructing continuous state trajectories of chaotic dynamical systems from sparse, noisy observations remains a fundamental open problem in nonlinear science. We introduce the Physics-Informed
Diffusion Model with Dormand-Prince Integration (PIDM-DP), which embeds a fully differentiable 5th-order Dormand-Prince (DP-RK45) ODE integrator directly into the reverse sampling loop of a Denoising
Diffusion Probabilistic Model (DDPM). At each denoising step, physics residuals are back-propagated via automatic differentiation, constraining every generated trajectory to satisfy the system's governing equations to 5th-order accuracy. A linear-scheduled guidance mechanism that ramps the physics weight from zero at high noise levels to its full value near the clean-data limit prevents the gradient explosions that cause naive physics-informed approaches to fail on stiff systems with Jacobian
eigenvalues of order $O(10^3)$. Evaluated across five
benchmark systems of increasing complexity 3D Lorenz, 3D Rössler, 5D Hyperchaotic, 20D Lorenz-96, and the stiff 3D Rabinovich-Fabrikant at 10% observation density with additive Gaussian noise ($\sigma=0.05$), PIDM-DP achieves reconstruction RMSE improvements of up to $15.4\times$ over an unconstrained
diffusion baseline and decisively outperforms the Ensemble Kalman Filter on stiff systems where ensemble covariance collapses. On the Rabinovich-Fabrikant out-of-distribution
benchmark, PIDM-DP attains RMSE $0.1097 \pm 0.0269$ versus $0.9443 \pm 0.5288$ (unconstrained
diffusion, $8.6\times$ worse) and $0.3561 \pm 0.3040$ (EnKF, $3.2\times$ worse), with $p<0.001$ in paired Wilcoxon tests ($N = 30$). Topological validation via the Rosenstein Lyapunov
estimator confirms that PIDM-DP preserves the chaotic invariant measure.
Read the full abstract ▾
Machine Learning
Spend Your Rollouts Where It Counts: Rollout Allocation for Group-Based RL Post-Training
Woojeong Kim, Ziyi Yang, Jing Nathan Yan et al. (+1 more)
Reinforcement learning (RL) is the dominant paradigm for post-training large
language models. However, in the online, on-policy setting, rollout generation dominates the computational cost of training. Group-based policy optimization methods compute advantages from multiple rollouts per
prompt, yet they indiscriminately allocate budget to
prompts with collapsed reward distributions, wasting expensive rollouts on negligible learning signals. We demonstrate that group-based updates are most effective in regimes of high reward variance. Since the policy evolves throughout training,
prompt informativeness must be estimated online rather than precomputed, but exhaustively evaluating every
prompt is computationally prohibitive. We introduce Pilot-Commit, a budget-aware rollout allocation framework for group-based RL post-training. Pilot-Commit decouples
prompt evaluation from exploitation: a pilot stage estimates per-
prompt informativeness using a fraction of the budget, and the remaining rollouts are allocated to high-leverage
prompts while low-signal
prompts are skipped. Across multiple math
reasoning benchmarks and model scales from 1.5B to 14B parameters, Pilot-Commit matches baseline accuracy with significantly lower sampling costs, reaching target accuracy up to $1.9\times$ faster than GRPO and $4.0\times$ faster than DAPO in cumulative rollouts.
Read the full abstract ▾
All authors Woojeong Kim, Ziyi Yang, Jing Nathan Yan, Jialu Liu
Machine Learning
Geometry-Aware Contrastive Learning for Few-Shot Automatic Modulation Recognition
Guanqun Zhao, Yitong Liu, Jiaxuan Fang et al. (+2 more)
Standard Self-Supervised Learning (SSL) for Automatic Modulation Recognition (AMR) struggles with ineffective isotropic augmentations, spectral instability, and semantic drift. To address these challenges, we propose Dynamic-Consistency Contrastive Learning (DyCo-CL), a geometry-aware framework that couples Virtual
Adversarial Augmentation (VAA) with a semantic consistency loss. We provide a theoretical analysis indicating that this strategy acts as an implicit spectral regularizer for the encoder, enabling stable
manifold exploration. Complementing this, our Signal-Adaptive Swin Backbone with fixed-window
attention improves structural stability by constraining
attention locality, while a Hybrid Knowledge Fusion module anchors representations with physical priors. Experiments on RML
benchmarks show that DyCo-CL achieves a 6.27% accuracy gain in 1-shot settings over prior methods.
Read the full abstract ▾
All authors Guanqun Zhao, Yitong Liu, Jiaxuan Fang, Yufei Mao, Hongwen Yang
Machine Learning · Cross-listed from cs.CR
Cordyceps: Covert Control Attacks on LLMs via Data Poisoning
Zedian Shao, Charles Fleming and Teodora Baluta
Large
language models (LLMs) are often fine-tuned on uncurated text
datasets that adversaries can poison. Existing poisoning attacks primarily rely on fixed trigger phrases that defenses such as outlier detection, clean-data regularization, or online monitoring can neutralize. In this paper, we propose a data poisoning method that teaches an LLM an information hiding scheme reliably and stealthily through semantic associations between shared knowledge such as facts or concepts and attacker-chosen phrases. The induced hiding scheme can encode and decode arbitrary malicious instructions, thus revealing a new and subtle poisoning-induced vulnerability: covert control attacks. We precisely characterize covert control attacks and evaluate them across $5$ LLMs, $3$
backdoor defenses, and $4$
prompt injection defenses. With a small poisoned fraction, covert control attacks outperform heuristic-based
prompt injection attacks in average attack success rate by about $40\%$ relative to clean fine-tuned models. They also circumvent defenses based on detection and fine-tuning, maintaining up to $93\%$ attack success rate after
backdoor defenses and up to $98\%$ after
prompt injection defenses.
Read the full abstract ▾