Machine Learning
A New Framework to Analyse the Distributional Robustness of Deep Neural Networks
Divij Khaitan and Subhashis Banerjee
Deep neural networks have achieved impressive performance on a variety of tasks, but their brittleness to distributional shifts remains a significant barrier to real-world deployment. In this paper, we propose a framework to analyse and quantify the distributional robustness of neural networks by studying the interactions between layer weights and activations. We model these interactions using Bernoulli distributions, using the separation between classes as a diagnostic proxy for robustness. We demonstrate the usefulness of this framework through models trained on CIFAR-10 and ImageNet. We show that our proposed metrics can distinguish between networks that have memorised their training data and those that have not. We also perform analogous experiments in the activation space and find that the same properties do not hold up. Additionally, we investigate the behaviour of our metrics under various distribution shifts and show that these shifts reduce separation under our path-based diagnostics. Our results suggest that this framework provides useful model-level diagnostics of representation structure and robustness.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.DC
Frontier: Towards Comprehensive and Accurate LLM Inference Simulation
Yicheng Feng, Xin Tan, Yangtao Deng et al. (+3 more)
Modern LLM serving is no longer homogeneous or monolithic. Production systems now combine disaggregated execution, complex parallelism, runtime optimizations, and stateful workloads such as
reasoning,
agents, and RL rollouts. Simulation is attractive for exploring this growing design space, yet existing simulators lack the architectural completeness and decision-grade fidelity it demands. Their monolithic-replica abstractions are ill-suited to disaggregated serving, while average-case analytical proxies can distort SLA predictions and even reverse optimization conclusions. We present Frontier, a discrete-event simulator for modern LLM inference serving. Frontier features a disaggregated abstraction. It captures the structure and dynamics of modern serving systems by modeling co-location, Prefill-Decode Disaggregation (PDD), and
Attention-FFN Disaggregation (AFD) with role-specific cluster workers, incorporating key runtime optimizations (e.g., CUDA Graphs, speculative decoding) within the scheduler-batch-engine loop, and supporting stateful requests for emerging workloads. It further provides accurate and generalizable predictions of computation, communication, and
memory costs across diverse serving scenarios with complex workload compositions. On 16-H800 GPU testbed, Frontier achieves an average throughput error below 4%. Compared with state-of-the-art simulators, it reduces end-to-end latency error from 44.9% to 6.4% under co-location and from 51.7% to 2.6% under disaggregation. It scales to over 1K GPUs on commodity CPUs and enables new use cases such as SLA-dependent Pareto frontier exploration, heterogeneous disaggregated allocation, agentic
reasoning scheduling validation, and RL post-training reconfiguration. We release Frontier at https://github.com/NetX-lab/Frontier.
Read the full abstract ▾
All authors Yicheng Feng, Xin Tan, Yangtao Deng, Yimin Jiang, Yibo Zhu, Hong Xu
Machine Learning
DeCoR: Design and Control Co-Optimization for Urban Streets Using Reinforcement Learning
Bibek Poudel, Lei Zhu, Kevin Heaslip et al. (+2 more)
Modern
vision systems can detect, track, and forecast urban actors at scale, yet translating perception outputs to urban design remains limited. We introduce DeCoR, a two-stage reinforcement learning framework that leverages flow observations to co-optimize crosswalk layout and network-level signal control. The design stage encodes the pedestrian network as a graph and learns a generative policy that parameterizes a Gaussian mixture model over crosswalk location and width, from which new crosswalks are sampled. For each layout, a shared control policy learns adaptive signal timings to minimize joint pedestrian and vehicle delay. On a 750 m real-world urban corridor with demand sensed from
video and Wi-Fi logs, DeCoR learns a layout that reduces pedestrian arrival time to their nearest crosswalk by 23% while using fewer crosswalks than existing configurations. On the control side, DeCoR reduces pedestrian and vehicle wait time by 79% and 65%, respectively, relative to fixed-time signalization. Further, the control policy generalizes to demands outside of training and is robust to layout changes without retraining.
Read the full abstract ▾
All authors Bibek Poudel, Lei Zhu, Kevin Heaslip, Sai Swaminathan, Weizi Li
Machine Learning
From Circuit Evidence to Mechanistic Theory: An Inductive Logic Approach
Nura Aljaafari, Danilo S. Carvalho and Andre Freitas
Mechanistic interpretability produces
circuit-level causal analyses of neural network behaviour, but discovered
circuits often remain isolated experimental artefacts: there is no shared formal representation for what
circuits compute, how they relate, or when two findings provide evidence for the same mechanism. This work provides a formal infrastructure for cumulative mechanistic science by treating
circuit interpretation as inductive theory construction. Each
circuit is characterised at two levels: a Causal Functional Signature (CFS), which grounds component behaviour in causal attribution evidence and
token role profiles, and an architectural signature $\tau_{\mathrm{arch}}$, learned by inductive logic programming (ILP) from scale-invariant structural predicates. Together, these constitute a formal coherence layer that makes mechanistic claims explicit, comparable via $\theta$-subsumption, and portable across model scales. CFS reveals qualitatively distinct computational strategies across task types, including
attention-mediated copying versus MLP-mediated binding. ILP signatures achieve substantially better structural separation than graph kernel and feature-vector baselines, and support principled transfer across model scales and architecture families.
Read the full abstract ▾
Machine Learning
Automatic Discovery of Disease Subgroups by Contrasting with Healthy Controls
Robin Louiset, Edouard Duchesnay, Benoit Dufumier et al. (+2 more)
In biomedical Subgroup Discovery, practitioners are interested in discovering interpretable and homogeneous subgroups within a group of patients. In this paper, assuming that healthy subjects (i.e., controls) share common but irrelevant factors of variation with the patients, we motivate and develop a Contrastive Subgroup Discovery method, entitled Deep UCSL. By contrasting patients with controls, Deep UCSL identifies subgroups driven solely by pathological factors, ignoring common variability shared with healthy subjects. Our framework employs a deep feature extractor to learn a discriminative representation space. Mathematically, we derive a novel loss based on the conditional joint likelihood of latent clusters and patient/control labels, optimized via an Expectation-Maximization strategy alternating between subgroup inference and feature encoder updates. A regularization term further encourages representations to capture disease-specific variability while ignoring variability shared with controls. Compared to previous related works, our approach quantitatively improves the quality of the estimated subgroups, as demonstrated on a MNIST example and four distinct real medical imaging
datasets. Code and
datasets are available at: https://github.com/rlouiset/deep_ucsl.
Read the full abstract ▾
All authors Robin Louiset, Edouard Duchesnay, Benoit Dufumier, Antoine Grigis, Pietro Gori
Machine Learning
TimeSRL: Generalizable Time-Series Behavioral Modeling via Semantic RL-Tuned LLMs – A Case Study in Mental Health
Yuang Fan, Lilin Xu, Millie Wu et al. (+8 more)
Longitudinal passive sensing enables continuous health prediction, yet models often fail under cross-
dataset distribution shifts. Traditional ML overfits cohort-specific artifacts, while Large
Language Models (LLMs) struggle to reason reliably over long, heterogeneous time-series. We introduce TimeSRL, a two-stage LLM framework that routes predictions through an explicit semantic bottleneck. The model first abstracts raw signals into high-level natural
language, then predicts behavioral outcomes from these abstractions alone. This forces the model to reason over semantic concepts that we argue generalize better than raw numbers. We optimize this process end-to-end using Group Relative Policy Optimization (GRPO) with Reinforcement Learning from Verifiable Rewards (RLVR), learning outcome-aligned abstractions without gold intermediate annotations. Instantiated on mental-health prediction, TimeSRL achieves state-of-the-art performance on a
benchmark designed to stress-test cross-cohort generalization under a rigorous leave-one-
dataset-out (LOSO) protocol, reducing mean absolute error (MAE) over strong non-LLM ML and LLM baselines by 3.1–10.1% and 9.5–44.1% for anxiety, and 3.2–9.6% and 27.4–57.6% for depression (all $p$s<0.05). TimeSRL significantly outperforms prior methods in cross-
benchmark transfer across different sensing pipelines, rivaling its own within-domain performance without target-domain fine-tuning. These results demonstrate that semantic abstractions are reusable and point to a new direction for generalizable behavior modeling via RL-tuned LLMs.
Read the full abstract ▾
All authors Yuang Fan, Lilin Xu, Millie Wu, Jingping Nie, Qingyu Chen, Yuzhe Yang, Zhuo Zhang, Xin Liu, Subigya Nepal, Xiaofan Jiang, Xuhai "Orson" Xu
Machine Learning · Cross-listed from stat.ML
Unbiased Gradients, Moving Stability Boundaries: Exact Mini-Batch Geometry in Linear Self-Attention
Krishnakumar Balasubramanian
Unbiased
stochastic gradients can match the full gradient in expectation while changing finite-step stability. We study this effect in one-layer linear self-
attention for in-context linear regression, where shared-mode mini-batch training reduces exactly to a random two-factor map. The full-batch map preserves an
elliptic region, whereas sampled
curvature and target correlation create batch-dependent stability boundaries. Their combined fluctuation is centered at the parameter-update level but induces a strictly outward drift of the full-batch boundary coordinate. We derive an exact one-step crossing criterion, adaptive concentration bounds, and conditional drift identities, and prove almost-sure escape in a balanced switching model. For independent long
prompts, mode leakage decays at the inverse square root of
prompt length. For finite-
token softmax
attention, we derive an explicit boundary-crossing law governed by an effective context length that shrinks as
attention scores concentrate. Experiments in ambient
attention objectives verify the reduced predictions, and four trained causal softmax
Transformers show larger mean
curvature shifts for smaller first-step batches, consistent with the predicted nonlinear outward bias.
Read the full abstract ▾
Machine Learning
A Mechanistic Study of Tabular Foundation Models
Marin Biloš, James T. Wilson, Anderson Schneider et al. (+1 more)
Tabular foundation models with different architectures converge in accuracy across a range of classification and regression tasks. This raises questions a leaderboard cannot answer: (i) whether the models execute the same in-context
algorithm, (ii) where row, column, and class-permutation invariances originate, and (iii) how robust they are under perturbations engineered against the inferred mechanism. We characterize all three. The model families realize qualitatively distinct similarity-based readouts: from an
attention-weighted vote over context labels to a class-conditional mean readout, each confirmed by causal intervention. We find that the representation collapse highlighted in prior work is not a practical concern for them. Each model's permutation invariances trace to specific positional parameters whose removal preserves accuracy and makes approximate invariance exact. Perturbations engineered against each readout reproduce predicted failure modes; hub and rank attacks isolate them from refit baselines. Together these results give a mechanistic account of contemporary tabular foundation models and identify which inductive biases govern both their accuracy and characteristic failures.
Read the full abstract ▾
All authors Marin Biloš, James T. Wilson, Anderson Schneider, Yuriy Nevmyvaka
Machine Learning
Stochastic MeanFlow Policies: One-Step Generative Control with Entropic Mirror Descent
Zeyuan Wang, Da Li, Yulin Chen et al. (+5 more)
Combining
entropy regularisation with mirror descent in online off-policy reinforcement learning requires expressive policies with efficient sampling and explicit
entropy control. Gaussian policies struggle to represent
multimodal targets, whereas generative policies often require iterative sampling or lack tractable
entropy estimates. We introduce
Stochastic MeanFlow Policies (SMFP), a one-step generative policy class that combines MeanFlow-based noise-to-action mappings with Gaussian reparameterisation. The latent mapping captures diverse action modes, while a learnable conditional Gaussian scale provides an analytic
entropy lower bound. We use this bound to regulate stochasticity alongside direct value guidance and advantage-weighted MeanFlow regression, enabling
entropy-regularised mirror-descent learning with one-step policies. Across seven MuJoCo
benchmarks, SMFP achieves strong performance against Gaussian and generative baselines with single-step inference. Additional experiments demonstrate
multimodal action modelling and effective exploration in sparse-reward
navigation.
Read the full abstract ▾
All authors Zeyuan Wang, Da Li, Yulin Chen, Yuehu Gong, Yanming Guo, Liang Bai, Tianyuan Yu, Yanwei Fu
Machine Learning
How Much Online RL is Enough? Informative Rollouts for Offline Preference Optimization in RLVR
Richa Verma and Balaraman Ravindran
Reinforcement Learning from Verifiable Rewards (RLVR) has emerged as a powerful paradigm for
reasoning in
language models, with GRPO as its primary example. However, GRPO requires continuous online rollout generation, making it computationally expensive and difficult to scale. While Direct Preference Optimization (DPO) offers a stable and efficient offline alternative, it is typically expected to underperform w.r.t. online RL methods such as GRPO when trained on rollouts from a cold supervised fine-tuned (SFT) policy. We introduce G2D (GRPO to DPO)}, a three-stage pipeline that performs a short GRPO warm-up, constructs a static preference
dataset, and fine-tunes a model offline with DPO. Across a set of values of the number of online steps (K) in GRPO on Qwen2.5-7B and Llama-3.1-8B, we find that offline DPO with moderate warm-up matches or outperforms GRPO at substantially lower compute cost in our setting. On Qwen2.5-7B, G2D at K=150 achieves 62.4% on MATH-500, outperforming GRPO (51.6%) by 10.8% at ~4x lower compute. On Llama-3.1-8B, G2D at K=500 achieves 49.4%, surpassing GRPO in our experimental setting. We show that performance is not governed by the number of preference pairs, which does not vary much w.r.t. K, but by their informativeness. Moderate warm-up produces rollouts with calibrated uncertainty, yielding stronger contrastive signal, while excessive warm-up leads to overconfident policies and less informative data. Our results recast the offline-online gap in RLVR as primarily a data informativeness problem, and identify short online RL warm-up with appropriate difficulty calibration of the fine-tuning
dataset as a compute-efficient alternative to online RL.
Read the full abstract ▾
Machine Learning
FedCoE: Bridging Generalization and Personalization via Federated Coordinated Dual-level MoEs
Penglin Dai, Fulian Li, Xincao Xu et al. (+3 more)
Federated Learning (FL) has emerged as a promising paradigm for
privacy-preserving distributed learning. However, existing FL methods face a fundamental challenge. Traditional averaging-based approaches suffer from parameter divergence under non-IID conditions, while personalized FL methods overfit to local data and fail to generalize to new clients (cold-start problem). Mixture-of-Experts naturally addresses this by routing heterogeneous data to specialized experts rather than forcing uniform aggregation. In this paper, we propose FedCoE, a
Federated Coordinated dual-level mixture-of-Experts framework that effectively balances global generalization with local personalization. FedCoE maintains multiple independent global expert models on the server and employs a shared gating network to dynamically model client-expert correlations during aggregation, effectively mitigating expert drift and gating inconsistency. To address the cold-start challenge, we introduce an adaptive mechanism that enables new clients to immediately leverage the global expert pool without extensive local training. Extensive experiments demonstrate that FedCoE achieves 78.00% global accuracy and 89.32% personalized accuracy on average, outperforming the baseline by 8.82% and 29.19%, respectively. In cold-start scenarios, FedCoE delivers 77.27% accuracy without any local fine-tuning, outperforming baselines by over 12.54%.
Read the full abstract ▾
All authors Penglin Dai, Fulian Li, Xincao Xu, Junhua Wang, Lixin Duan, Xiao Wu
Machine Learning
Nonparametric Learning and Earning with One-Point Feedback under Nonstationarity
Xiangyu Yang, Feng Xu, Jian-Qiang Hu et al. (+1 more)
Firms increasingly rely on dynamic pricing to respond to evolving customer demand, yet in many applications they observe only the revenue generated by a single posted price in each period. At the same time,
market conditions may shift gradually or abruptly due to changes in customer preferences, competition, or external shocks. These features create two intertwined challenges: learning the revenue–demand relationship from limited feedback and adapting pricing decisions to a changing environment. We study how a seller can learn and earn effectively under these constraints, without assuming a specific parametric form for demand. We develop a learning framework that updates prices using revenue-based gradient approximations constructed from one observation per period. To address environmental changes, we incorporate a restarting mechanism that periodically refreshes the learning process so that outdated information is discounted. When the degree of nonstationarity is unknown, we further introduce a meta-learning layer to adaptively hedge across multiple restarting schedules. We provide performance guarantees for our approach, showing how cumulative revenue loss relative to a fully informed
benchmark depends on both the time horizon and the magnitude of
market variation. Simulation experiments using synthetic and real-world data illustrate the effectiveness of the proposed procedures.
Read the full abstract ▾
All authors Xiangyu Yang, Feng Xu, Jian-Qiang Hu, Jiaqiao Hu
Machine Learning
On the Cost and Benefit of Chain of Thought: A Learning-Theoretic Perspective
Yue Zhang, Zhiyi Dong, Tommaso Cesari et al. (+1 more)
We develop a learning-theoretic framework for understanding Chain of Thought (CoT). We model CoT as the interaction between an answer map and a chain rule that generates intermediate questions autoregressively, and define the
reasoning risk of a hypothesis under this interaction. Our first result is a tight canonical decomposition of this risk into two terms with opposing roles: an oracle-trajectory risk (OTR), which captures the benefit of CoT and reduces to a target-domain risk in a domain adaptation problem, and a trajectory-mismatch risk (TMR), which captures the cost of CoT through error accumulation along mismatched
reasoning trajectories. We then show that this cost is unavoidable without structure: if any one of the loss, the hypothesis answer map, or the chain rule lacks stability, the TMR can be arbitrarily large even when the OTR is zero and the hypothesis is uniformly close to the ground truth. Conversely, under stability, we prove a tight upper bound on the TMR governed by an exact amplification factor that identifies bounded, linear, and exponential error-growth regimes. Together, these results give a precise theory of when CoT helps, when it hurts, and what controls the transition between the two.
Read the full abstract ▾
All authors Yue Zhang, Zhiyi Dong, Tommaso Cesari, Yongyi Mao
Machine Learning · Cross-listed from stat.ML
Theoretical guidelines for annealed Langevin dynamics in compositional simulation-based inference
Camille Touron, Gabriel V. Cardoso, Julyan Arbel et al. (+1 more)
Compositional score-based approaches to simulation-based inference (SBI) approximate the posterior over a shared parameter given $n$ independent observations by aggregating individually learned posterior scores: currently, there are two main propositions of such methods (Geffner et al. (2023), Linhart et al. (2026)). As the resulting composite score does not correspond to the score of any distribution along the forward
diffusion path of the true multi-observation posterior, sampling from it via a reverse SDE leads to an irreducible bias. Annealed Langevin dynamics provides a principled alternative: it treats the composite score as the genuine score of a sequence of tractable bridging densities and samples from them in succession. When properly tuned, it could lead to a controllable bias. However, its hyperparameters, namely step sizes, the number of steps per level, and the number of annealing levels, have so far been chosen empirically. We derive Wasserstein bounds for annealed Langevin with approximate scores and translate them into explicit decision rules for these hyperparameters that guarantee a prescribed sampling accuracy, while highlighting different theoretical aspects of each composite score formulation. In the Gaussian setting, we obtain closed-form expressions for all relevant quantities and prove that the bridging densities of Linhart et al. (2026) consistently admit larger step sizes and require fewer total Langevin steps than those of Geffner et al. (2023). Furthermore, we show empirically that the tuning obtained in the Gaussian setting generalizes to more complex problems, thus providing a well-understood and theoretically grounded starting point for practitioners using compositional score-based approaches.
Read the full abstract ▾
All authors Camille Touron, Gabriel V. Cardoso, Julyan Arbel, Pedro L. C. Rodrigues
Machine Learning
Graph Navier Stokes Networks
Zexing Zhao, Guangsi Shi, Yu Gong et al. (+4 more)
Graph Neural Networks (GNNs) have emerged as a cornerstone of deep learning, with most existing methods rooted in graph signal processing and
diffusion equations to model message passing. However, these approaches inherently suffer from the oversmoothing problem, where node features become indistinguishable as the network depth increases. Inspired by the Navier Stokes equations, we introduce Graph Navier Stokes Networks (GNSN), a novel architecture that transcends conventional
diffusion-based message passing by incorporating convection into graph structures. GNSN defines a dynamic velocity field on the graph to govern convection, enabling more efficient and direct message propagation. By adaptively balancing convection and
diffusion, GNSN is able to efficiently handle
datasets with varying levels of homophily. Extensive evaluations across twelve real-world
datasets demonstrate that GNSN consistently outperforms state-of-the-art baselines in classification accuracy. Moreover, experimental results further emphasize its effectiveness in alleviating the oversmoothing problem.
Read the full abstract ▾
All authors Zexing Zhao, Guangsi Shi, Yu Gong, Tianyu Wang, Shirui Pan, Hongye Cheng, Yuxiao Li
Machine Learning
Divide and Contrast: Learning Robust Temporal Features without Augmentation
Abdul-Kazeem Shamba, Kerstin Bach and Gavin Taylor
Self-supervised learning for time-series representation aims to reduce reliance on labeled data while maintaining strong downstream performance, yet many existing approaches incur high computational costs or rely on assumptions that do not hold across diverse temporal dynamics. In this work, we introduce Divide and Contrast (Di-COT), an unsupervised framework that avoids data augmentation and multiple encoder passes by contrasting informative substructures within a window rather than individual timesteps. Di-COT stochastically partitions each window into a small number of overlapping sub-blocks per iteration, enabling efficient and meaningful contrast while mitigating false positives during temporal transitions. To further improve scalability, we adopt a contrastive objective whose computation depends on the batch size and the number of sub-blocks, making loss computation independent of sequence length. Extensive experiments on six large-scale real-world
datasets, as well as the UCR and UEA
benchmarks, demonstrate that Di-COT learns semantically structured and transferable representations, achieving state-of-the-art performance on classification, clustering, $k$NN, and cross-
dataset transfer, while substantially reducing training time. The source code is publicly available at https://github.com/sfi-norwai/Di-COT.
Read the full abstract ▾
Machine Learning
APEX: Autonomous Policy Exploration for Self-Evolving LLM Agents
Yibo Li, Jiashuo Yang, Zhi Zheng et al. (+5 more)
LLM
agents have shown strong performance across a wide range of complex tasks, including interactive environments that require long-horizon decision making. But these
agents cannot learn on the fly at test time. Self-evolving
agents address this by accumulating
memory and reflection across episodes rather than requiring model-weight updates. However, these
agents often suffer from exploration collapse: as
memory grows, behavior concentrates around familiar high-reward routines, reducing the chance of discovering better alternatives. To address this problem, we propose Autonomous Policy EXploration (APEX), which builds and maintains an explicit strategy space through a strategy map-a directed acyclic graph of milestones with prerequisite dependency edges. In APEX, Fork Discovery expands the map with evidence-grounded unexplored directions, while Policy Selection balances exploration and exploitation during planning. Evaluated on nine Jericho text-adventure games and WebArena, a realistic web interaction
benchmark, APEX outperforms all baselines. Extensive ablations validate each component's contribution and demonstrate robustness across diverse settings, demonstrating APEX's effectiveness for sustained exploration in self-evolving
agents.
Read the full abstract ▾
All authors Yibo Li, Jiashuo Yang, Zhi Zheng, Zhiyuan Hu, Yuan Sui, Shizun Wang, Yufei He, Bryan Hooi
Machine Learning
OCTOPUS: Optimized KV Cache for Transformers via Octahedral Parametrization Under optimal Squared error quantization
Mark Boss, Vikram Voleti, Simon Donné et al. (+1 more)
The key-value (KV) cache dominates
memory bandwidth and footprint in long-context autoregressive inference. Recent rotation-preconditioned codecs (TurboQuant, PolarQuant) show that a structured random rotation followed by a per-coordinate scalar quantizer matched to an analytically tractable marginal is a near-optimal recipe for KV compression. OCTOPUS advances this paradigm through joint quantization of rotated coordinate triplets. Each triplet's direction is mapped to a square via an octahedral parameterization, and the two resulting coordinates and the triplet norm are Lloyd-Max quantized against implementation-matched marginals. Optimizing the per-triplet squared error gives a strictly non-uniform bit allocation depending only on the total dimensionality of the keys. We find the finite-dimensional quality optimum with sweeps to be constant on every real decoder we test. The codec is data-oblivious, online, and deterministic given a seed. Across text,
video, and audio, OCTOPUS matches or beats every prior rotation codec at every reported bit width and metric, with a lead that grows as bits drop for extreme compression. Furthermore, a fused Triton implementation reconstructs keys on the fly without materializing the uncompressed key, so the codec adds no decode-time bandwidth or latency over the existing dequantization. Project Page: https://octopus-quant.github.io/
Read the full abstract ▾
All authors Mark Boss, Vikram Voleti, Simon Donné, Shimon Vainer
Machine Learning
PREFINE: Preference-Based Implicit Reward and Cost Fine-Tuning for Safety Alignment
Richa Verma, Bavish Kulur, Sanjay Chawla et al. (+1 more)
We address the problem of making a pre-trained reinforcement learning (RL) policy safety-aware by incorporating cost constraints without retraining it from scratch. While costs could be numerically encoded, we assume a more general setting is when costs are provided as preferences. Given a reward-optimized policy and a small
dataset of preferred (low-cost) and dispreferred (high-cost) trajectories, our goal is to fine-tune the policy to generate low-cost behaviors while retaining high rewards. Unlike standard RLHF in
language models, where preferences are defined over responses to the same
prompt, our setting involves trajectory-level preferences in continuous control environments. We introduce PREFINE: Preference-based Implicit Reward and Cost Fine-Tuning for Safety
Alignment which is a preference-based fine-tuning method that adapts Direct Preference Optimization (DPO), which is now widely used for LLM fine-tuning, to the sequential decision making setting. PREFINE constructs policy-sampled counterfactual trajectories to establish meaningful preference contrasts and jointly optimizes for reward retention and safety
alignment. Empirically, PREFINE reduces constraint violations and catastrophic failures by over 60% while maintaining original reward behavior. PREFINE produces policies that achieve low-cost, high-reward performance with significantly improved data and computational efficiency compared to full offline RL or imitation learning, bridging preference
alignment and safe policy adaptation in continuous domains.
Read the full abstract ▾
All authors Richa Verma, Bavish Kulur, Sanjay Chawla, Balaraman Ravindran
Machine Learning · Cross-listed from stat.ML
Federated LoRA Fine-Tuning for LLMs via Collaborative Alignment
Shuaida He, Liwen Chen and Long Feng
Low-rank adaptation (LoRA) has emerged as a powerful tool for parameter-efficient fine-tuning of large
language models (LLMs). This paper studies LoRA under a
federated learning setting, enabling collaborative fine-tuning across clients while preserving parameter efficiency. We focus on a highly heterogeneous regime in which clients share only partial structure and a substantial subset may be contaminated. We propose Collaborative Low-rank
Alignment and Identifiable Recovery (CLAIR), a contamination-aware framework that relies only on preliminary local
estimators. Its formulation applies broadly, from linear regression to neural network and LLM modules, whenever local adaptation can be represented by
matrix-valued updates. CLAIR recovers the shared LoRA subspace and detects contaminated clients via a structured low-rank plus block-sparse decomposition. We prove exact recovery of the shared LoRA subspace in the noiseless case, stable recovery under preliminary estimation error, and consistent collaborative-set recovery under mild separation conditions. We further quantify the gain from CLAIR refinement: it reduces off-subspace estimation error through cross-client averaging while preserving client-specific variation within the shared LoRA subspace, thus improves over local fine-tuning whenever this oracle gain outweighs the costs of subspace estimation and benign-client heterogeneity. Empirically, we demonstrate the benefits of CLAIR by fine-tuning a
Transformer architecture on a text-copying task. The results show accurate contamination detection and improved benign-client performance compared with local fine-tuning and non-robust
federated averaging.
Read the full abstract ▾
Machine Learning
Behavior-Consistent Deep Reinforcement Learning
Marcel Hussing, Liv G. d'Aliberti, Claas Voelcker et al. (+2 more)
Reinforcement learning (RL) often exhibits high variance across training runs, leading to unreliable performance and posing a major challenge to deployment in real-world domains. In this work, we address the challenge of cross-run policy divergence by formalizing the problem of behavior-consistent RL, where the objective is to obtain policies that are both high-performing and distributionally similar across training runs. Our key observation is that maximum-
entropy RL provides a direct mechanism for controlling behavioral divergence by anchoring runs to a common (uniform) prior. We prove that, for Boltzmann policies, choosing the temperature proportional to $Q$-function disagreement bounds the pairwise KL divergence between the induced policies. However, we also show that naïvely increasing
entropy might impair policy optimization while amplifying off-policy error. Building upon these observations, we propose $Q$-value Expectile Disagreement (QED), a state-dependent temperature schedule that uses double-critic disagreement as a single-run proxy for cross-run disagreement. Empirically, we demonstrate that across 18 continuous-control tasks, QED reduces across-run divergence by two orders of magnitude without sacrificing performance, resulting in a considerable reduction in return variance at modest sample-efficiency costs.
Read the full abstract ▾
All authors Marcel Hussing, Liv G. d'Aliberti, Claas Voelcker, Benjamin Eysenbach, Eric Eaton
Machine Learning · Cross-listed from quant-ph
Enhanced Reinforcement Learning-based Process Synthesis via Quantum Computing
Austin Braniff, Fengqi You and Yuhe Tian
In this work, we present
quantum reinforcement learning (RL) as a solution strategy for process synthesis problems. Building on our prior work, we develop a generalized framework that formally poses process synthesis as a Markov decision process and introduces
quantum-enhanced RL
algorithms to solve it with improved scalability. Earlier implementations of
quantum-based RL for process synthesis were limited by
qubit requirements, which scaled poorly with problem complexity. This work overcomes this challenge by introducing state encoding
algorithms to decouple
qubit requirements from problem size. A classical RL-based solution strategy is used as a baseline to
benchmark the
quantum algorithms under identical training conditions. All
algorithms are evaluated across a flowsheet synthesis problem of increasing unit counts to analyze their performance and scalability. Results show that all approaches are capable of identifying the optimal flowsheet designs in small design spaces. For moderate-scale unit counts,
quantum approaches demonstrate competitive performance on a per-episode basis and improved efficiency on a per-parameter basis versus the classical RL
benchmark. This work provides a foundation for future
quantum computing applications within process systems engineering, establishes a controlled
benchmark for comparing classical and
quantum algorithms, and shows that the proposed
quantum variants remain competitive for the process synthesis problem examined in this work.
Read the full abstract ▾
Machine Learning · Cross-listed from eess.SY
Reinforcement Learning-based Control via Y-wise Affine Neural Networks: Comparative Case Studies for Chemical Processes
Austin Braniff and Yuhe Tian
In this work we present an efficient and practically implementable approach for the application of reinforcement learning (RL)-based control in chemical process systems. This is an area that has yet to widely adopt RL-based control largely due to inherent challenges in trusting RL
algorithms and the time-consuming process of training reliable
agents. To address these challenges, we leverage a class of RL
algorithms termed Y-wise Affine Neural Network (YANN)- RL, which we have developed in our prior work (Braniff and Tian, 2025a). By strategically initializing actor and critic networks YANN-RL
algorithms provide confident and interpretable starting points within control schemes. We apply this RL-based control approach to three different process engineering case studies publicly available on the PC-Gym library (Bloor et al., 2026): (i) a continuous stirred tank reactor (CSTR), (ii) a four-tank system, and (iii) a multistage extraction column. Our approach is compared to several popular RL
algorithms (PPO, SAC, DDPG, and TD3) and is benchmarked against nonlinear model predictive control (NMPC). These case studies demonstrate that YANN-RL can greatly reduce the training time and data needed, can be deployed with confidence for chemical process systems, and can approach the performance of NMPC without the knowledge of a full nonlinear model.
Read the full abstract ▾
Machine Learning
Domain-Adaptable Reinforcement Learning for Code Generation with Dense Rewards
Erfan Aghadavoodi Jolfaei, Daniel Maninger, Abhinav Anand et al. (+2 more)
Large
language models show strong potential for automated code generation, but lack guarantees for correctness, quality, safety, and domain-specific constraints. For instance in robotics, where code generation is increasingly being used for planning and executing actions, awareness of the environment and physical constraints is critical. To facilitate the adaption of code-generating LLMs to diverse requirements, including domain-specific ones, we present a reinforcement learning framework that fine-tunes pre-trained LLMs using proximal policy optimization. Our customizable execution-aware reward formula captures and optimizes syntax, functional correctness, code style, security, and simulator executability. A
token-level reward mapping mechanism enables effective credit assignment from execution outcomes to generated
tokens. The framework is evaluated on general-purpose code generation (MBPP/MBPP+) and robotic program synthesis (RoboEval). The results show substantial improvements in functional correctness and simulator executability, including an absolute pass@1 increase of 19% on MBPP and a reduction in execution failures by 51% on RoboEval. These findings demonstrate that structured reinforcement learning can effectively align
language models to correct program generation and domain-specific requirements.
Read the full abstract ▾
All authors Erfan Aghadavoodi Jolfaei, Daniel Maninger, Abhinav Anand, Mert Tiftikci, Mira Mezini
Machine Learning
ChunkFT: Byte-Streamed Optimization for Memory-Efficient Full Fine-Tuning
Yongkang Liu, Zijing Wang, Mengjie Zhao et al. (+7 more)
This work presents ChunkFT, a
memory-efficient fine-tuning framework that reformulates full-parameter fine-tuning around a dynamically activated working set. ChunkFT enables gradient computation for arbitrary sub-tensors without modifying the network architecture, providing an algorithmic foundation for optimizing arbitrary sub-networks while avoiding standard dense gradient computation. We provide a theoretical convergence analysis of ChunkFT in the deterministic setting. Empirically, we apply ChunkFT to fine-tune Llama 3-8B and Llama 3-70B using a single RTX 4090-24GB GPU and 2$\times$ H800-80GB GPUs, respectively. Full-parameter fine-tuning of a 7B model with a 1K input length requires only 13.72GB of GPU
memory. The results demonstrate the effectiveness of ChunkFT in
memory usage, running time, and optimization quality. Moreover, downstream evaluations on
language understanding, mathematical
reasoning, and MT-Bench show that ChunkFT consistently outperforms existing
memory-efficient baselines. Notably, ChunkFT achieves performance comparable to, and in some cases exceeding, full-parameter fine-tuning. Our repository is on https://github.com/misonsky/chunk.
Read the full abstract ▾
All authors Yongkang Liu, Zijing Wang, Mengjie Zhao, Ercong Nie, Mingyang Wang, Qian Li, Feiliang Ren, Shi Feng, Daling Wang, Hinrich Schütze
Machine Learning · Cross-listed from stat.ML
A Rigorous, Tractable Measure of Model Complexity
Oskar Allerbo and Thomas B. Schön
An accurate assessment of a model's complexity is crucial for topics such as interpretation, generalization, and model selection. However, existing complexity measures either rely on heuristic assumptions or are computationally prohibitive. In this paper, we present a mathematically rigorous yet easy-to-compute measure of model complexity that is based on the similarities between feature vectors for different inputs. It is thus well-defined for any parametric model (via NTK feature generalization) and for kernel-based non-parametric models (via the induced kernel). We prove that our measure of complexity generalizes model-specific complexity measures, including
polynomial degree (for
polynomial regression), kernel length scale (for Matérn kernels), number of neighbors (for k-nearest neighbors), number of splits (for decision trees), and number of trees (for random forests). We also use our measure to obtain new insights into the double descent phenomenon for random
Fourier features, random forests, neural networks, and gradient boosting.
Read the full abstract ▾
Machine Learning
Q-SYNTH: Hybrid Quantum-Classical Adversarial Augmentation for Imbalanced Fraud Detection
Adam Innan, Mansour El Alami, Nouhaila Innan et al. (+2 more)
Credit card fraud detection is fundamentally challenged by extreme class imbalance, where fraudulent transactions are rare yet operationally critical. This imbalance often biases supervised learners toward the legitimate class, leading to high overall accuracy but weaker fraud-class recall and F1-score. This paper introduces Q-SYNTH, a hybrid classical–
quantum generative
adversarial framework in which a parameterized
quantum circuit serves as the generator and a classical neural network serves as the discriminator. Q-SYNTH is designed for minority-class fraud synthesis in tabular data and is evaluated along two dimensions: statistical fidelity to real fraud samples and downstream performance for fraud detection. To this end, generated samples are assessed using distributional similarity measures based on Kolmogorov-Smirnov statistics and Wasserstein distances, real-vs-synthetic detectability measured by AUC-ROC, and downstream classification performance across both
quantum and classical classifiers. Under the reported protocol, Q-SYNTH reduces marginal distribution mismatch relative to a classical GAN baseline while maintaining competitive downstream fraud-detection performance. Although SMOTE achieves the strongest feature-wise similarity and the classical GAN attains the highest downstream performance in several settings, Q-SYNTH offers a favorable compromise between distributional fidelity and downstream performance, supporting the feasibility of hybrid
quantum augmentation for imbalanced fraud detection.
Read the full abstract ▾
All authors Adam Innan, Mansour El Alami, Nouhaila Innan, Muhammad Shafique, Mohamed Bennai
Machine Learning
Learning First Integrals via Backward-Generated Data and Guided Reinforcement Learning
Jingfeng Zhong, Zhengxiang Liu, Zhijie Wang et al. (+1 more)
The discovery of first integrals is of fundamental scientific importance for understanding conservation laws in dynamical systems. However, existing symbolic computation tools and Large
Language Models (LLMs) remain limited on this task because high-quality training data are scarce and successful solutions often depend on mathematical intuition. This paper presents FISolver, an LLM-based solver developed to address this challenge. First, we introduce a "Backward Generation"
algorithm that systematically builds large-scale
datasets of (differential equation, first integral) pairs by deriving differential equations from sampled integrals, thereby alleviating the data scarcity bottleneck. Second, we apply supervised fine-tuning to a compact mathematical model and further improve its performance through reinforcement learning with a Levenshtein Distance-based shaped reward. In addition, we design data synthesis and blending strategies that support effective adaptation to difficult problem families from sparse examples. Experiments show that FISolver, while requiring substantially lower computational cost, significantly outperforms larger mathematical LLMs and commercial solvers such as Mathematica on challenging
benchmarks, indicating a new data-driven route for automated discovery of first integrals.
Read the full abstract ▾
All authors Jingfeng Zhong, Zhengxiang Liu, Zhijie Wang, Shuai Li
Machine Learning · Cross-listed from cs.CV
Comparative Analysis of Military Detection Using Drone Imagery Across Multiple Visual Spectrums
Sourov Roy Shuvo, Prajwal Panth, Rajesh Chowdhury et al. (+3 more)
In modern warfare, drones are becoming an essential part of intelligence gathering and carrying out precise attacks in different kinds of hostile environments. Their ability to operate in real-time and hostile environments from a safe distance makes them invaluable for surveillance and military operations. The KIIT-MiTA
dataset is comprised of images of different military scenarios taken from drones, and these provide a foundation for detecting military objects, but it does not take into account the various types of real-world scenarios. With that in mind, to evaluate how the models are performing under varying conditions, four different types of
datasets are created: Gray Scale, Thermal
Vision, Night
Vision, and Obscura
Vision. These simulate the real-world environments such as low visibility, heat-based imagery, and nighttime conditions. The YOLOv11-small model is trained and used to detect objects across diverse settings. This research boosts the performance and reliability of drone-based operations by contributing to the development of advanced detection systems in both defensive and offensive missions.
Read the full abstract ▾
All authors Sourov Roy Shuvo, Prajwal Panth, Rajesh Chowdhury, Sorup Chakraborty, Sudip Chakrabarty, Prasant Kumar Pattnaik
Machine Learning · Cross-listed from cs.CL
Automated ICD Classification of Psychiatric Diagnoses: From Classical NLP to Large Language Models
Fernando Ortega, Raúl Lara-Cabrera, Jorge Dueñas-Lerín et al. (+3 more)
Mental health has become a global priority, leading to a massive administrative burden in the coding of clinical diagnoses. This study proposes the automation of psychiatric diagnostic analysis by mapping free-text descriptions to the International Classification of Diseases (ICD) using Natural
Language Processing (NLP) and Machine Learning (ML) techniques. Utilizing a specialized
dataset of 145,513 Spanish psychiatric descriptions, various text representation paradigms were evaluated, ranging from classical frequency-based models (BoW, TF-IDF) to state-of-the-art Large
Language Models (LLMs) such as e5\_large, BioLORD, and Llama-3-8B. Results indicate that
transformer-based
embeddings consistently outperform traditional methods by capturing implicit semantic cues and nuanced medical terminology. The e5\_large model, through end-to-end fine-tuning, achieved the highest performance with a $F1_{micro}$ score of 0.866. This research demonstrates that adapting LLMs to specific clinical nomenclature is essential for overcoming the challenges of “long-tail” label distributions and the inherent ambiguity of psychiatric discourse.
Read the full abstract ▾
All authors Fernando Ortega, Raúl Lara-Cabrera, Jorge Dueñas-Lerín, Alejandro de la Torre-Luque, Mercé Salvador Robert, Enrique Baca-García