Machine Learning
KL for a KL: On-Policy Distillation with Control Variate Baseline
Minjae Oh, Sangjun Song, Gyubin Choi et al. (+2 more)
On-Policy
Distillation (OPD) has emerged as a dominant post-training paradigm for Large
Language Models, especially in
reasoning domains. However, OPD remains unstable in practice due to the high gradient variance of its single-sample Monte Carlo
estimator, and recipes for stable training are still immature. We propose vOPD (On-Policy
Distillation with a control variate baseline), which casts OPD as policy-gradient RL and stabilizes it by introducing a control variate baseline—canonically a value function—from the RL literature. We show that the OPD value function admits a closed form: the per-
token negative reverse KL divergence between student and teacher, available directly from the existing forward pass with no additional critic or inference. Existing stabilization methods either compute the full
token-level reverse KL over the entire vocabulary, adding significant overhead, or restrict it to a top-k support, biasing the objective. vOPD instead preserves the lightweight single-sample
estimator and subtracts the value function as a detached baseline to keep the gradient unbiased while reducing variance. Moreover, we show that a top-k approximation of the baseline further lowers cost without compromising performance. Across mathematics, science, and instruction-following
benchmarks, vOPD consistently outperforms base OPD and matches the most expensive full-vocabulary method, offering an efficient stabilization of OPD through principled RL variance reduction.
Read the full abstract ▾
All authors Minjae Oh, Sangjun Song, Gyubin Choi, Yunho Choi, Yohan Jo
Machine Learning
ADKO: Agentic Decentralized Knowledge Optimization
Lucas Nerone Rillo, Zhanhong Jiang, Nastaran Saadati et al. (+4 more)
We present Agentic Decentralized Knowledge Optimization (ADKO), a framework for collaborative black-box optimization across autonomous
agents that achieves sample efficiency,
privacy preservation, heterogeneous-objective handling, and communication efficiency. Each
agent maintains a private Gaussian Process (GP) surrogate trained on local data and communicates only through knowledge
tokens-compact, lossy summaries containing directional signals, advantage scores, and optional
language-model (LM) insights-without sharing raw data or model parameters. ADKO unifies GP-Upper Confidence Bound (GP-UCB), parallel
Bayesian optimization, decentralized learning, and LM-guided discovery. We provide the first formal analysis of dual information loss:
token compression, quantified via mutual-information-based fidelity, and LM approximation error, decomposed into bias and
stochastic noise. Our main result shows cumulative regret decomposes into GP error, LM bias, LM noise, and compression loss, with necessary and sufficient conditions for sublinear regret. We also propose fidelity-aware
token pruning to preserve high-information
tokens under
memory budget. Experiments on neural architecture search and scientific discovery validate the theory and show consistent improvements over strong baselines.
Read the full abstract ▾
All authors Lucas Nerone Rillo, Zhanhong Jiang, Nastaran Saadati, Aditya Balu, Baskar Ganapathysubramanian, Chinmay Hegde, Soumik Sarkar
Machine Learning
On the Tradeoffs of On-Device Generative Models in Federated Predictive Maintenance Systems
Usevalad Milasheuski, Piero Baraldi, Enrico Zio et al. (+1 more)
Federated Learning (FL) has emerged as a promising paradigm for preserving client data ownership and control over distributed Internet of Things (IoT) environments. While discriminative models dominate most FL use cases, recent advances in generative models - such as Variational Autoencoders (VAE), Generative
Adversarial Networks (GAN), and
Diffusion Models (DM) - offer new opportunities for unsupervised anomaly detection in time series analysis, with relevant applications in predictive maintenance (PdM) in critical industrial infrastructures. In this work, we present a comprehensive analysis of VAEs, GANs, and DMs in the context of
federated PdM. We analyze their performance and communication overhead under both full and partial federation setups, where only subsets of model components are shared. Building on split-parameter personalized FL, this paper provides an empirical instantiation of shared/private parameter partitioning for unsupervised generative Time-Series Anomaly Detection (TSAD), mapping the split to analysis and synthesis modules across the generative architectures. Our experiments over five clients and two industrial time-series
benchmarks reveal distinct trade-offs in model utility, stability, and communication efficiency in heterogeneous cross-silo FL. For the evaluated GAN-based configurations, full federation improves final anomaly-detection utility relative to independent local training, although these models remain less robust than the VAE- and DDPM-based alternatives. For DMs, decoder sharing matches full federation while substantially reducing the communication payload.
Read the full abstract ▾
All authors Usevalad Milasheuski, Piero Baraldi, Enrico Zio, Stefano Savazzi
Machine Learning
Actor-Critic Algorithm for Dynamic Expectile and CVaR
Yudong Luo and Erick Delage
Optimizing dynamic risk with
stochastic policies is challenging in both policy updates and value learning. The former typically requires transition perturbation, while the latter may rely on model-based approaches. To address these challenges, we propose a surrogate policy gradient without transition perturbation under softmax policy parameterization. We further develop model-free value learning methods for dynamic expectile and conditional value-at-risk by leveraging elicitability. Finally, inspired by Expected SARSA and Expected Policy Gradient, a model-free off-policy actor-critic
algorithm is constructed. Empirical results in domains with verifiable risk-averse behavior show that our
algorithm can learn risk-averse policy and consistently outperforms other existing methods.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
MatryoshkaLoRA: Learning Accurate Hierarchical Low-Rank Representations for LLM Fine-Tuning
Ionut-Vlad Modoranu, Mher Safaryan and Dan Alistarh
With the rise in scale for deep learning models to billions of parameters, the computational cost of fine-tuning remains a significant barrier to deployment. While Low-Rank Adaptation (LoRA) has become the standard for parameter-efficient fine-tuning, the need to set a predefined, static rank $r$ requires exhaustive grid searches to balance efficiency and performance. Existing rank-adaptive solutions such as DyLoRA mitigate this by sampling ranks during the training from a predefined distribution. However, they often yield sub-optimal results at higher ranks due to lack of consistent gradient signals across the full hierarchy of ranks, thus making these methods data-inefficient. In this paper, we propose MatryoshkaLoRA, a general, Matryoshka-inspired training framework for LoRA that learns accurate hierarchical low-rank representations by inserting a fixed, carefully crafted diagonal
matrix $P$ between the existing LoRA adapters to scale their sub-ranks accordingly. By introducing this simple modification, our general framework recovers LoRA and DyLoRA only by changing $P$ and ensures all sub-ranks embed the available gradient information efficiently. Our MatryoshkaLoRA supports dynamic rank selection with minimal degradation in accuracy. We further propose Area Under the Rank Accuracy Curve (AURAC), a metric that consistently evaluates the performance of hierarchical low-rank adapters. Our results demonstrate that MatryoshkaLoRA learns more accurate hierarchical low-rank representations than prior rank-adaptive approaches and achieves superior accuracy-performance trade-offs across ranks on the evaluated
datasets. Our code is available at https://github.com/IST-DASLab/MatryoshkaLoRA.
Read the full abstract ▾
Machine Learning
Distributional simplicity bias and effective convexity in Energy Based Models
Aurélien Decelle, Alfonso de Jesús Navas Gómez and Beatriz Seoane
Energy-based learning is a powerful framework for generative modelling, but its training is inherently non-convex, leading potentially to sensitivity to initialisation, poor local optima, and unstable gradient dynamics. We present a dynamical analysis of energy-based learning through the lens of the effective model, which can be interpreted as either a generalised Ising model with higher-order interactions or the
Fourier expansion of the energy. Under sufficient expressivity, we show that the gradient flow induced by learning strictly positive distributions over binary variables admits two types of fixed points: data-consistent points, which exactly reproduce the target distribution, and spurious points, which satisfy stationarity without matching the target distribution. Around data-consistent points, we show that perturbations are either stable or neutral, with neutral directions leaving the effective model invariant. Finally, we show that gradient dynamics induce a hierarchy in which lower-order interactions are learned before higher-order ones. This provides a mechanistic explanation for the distributional simplicity bias and clarifies why fixed points that are not data-consistent at low orders are not observed in practice.
Read the full abstract ▾
Machine Learning
\mathsf{VISTA}: Decentralized Machine Learning in Adversary Dominated Environments
Hanzaleh Akbari Nodehi, Parsa Moradi, Soheil Mohajer et al. (+1 more)
Decentralized machine learning often relies on outsourcing computations, such as gradient evaluations, to untrusted worker nodes. Existing robust aggregation methods can mitigate malicious behavior under honest-majority assumptions, but may fail when adversaries control a majority of the workers. We study this adversary-dominated setting through an incentive-oriented framework in which reports are accepted and rewarded only when they are mutually consistent up to a threshold. This turns the adversary from a pure saboteur into a rational
agent that trades off increasing estimation error against the risk of rejection and loss of reward. We consider iterative optimization under this model. Unlike one-shot computation, iterative learning requires long-horizon decisions: permissive acceptance rules enable faster early progress but admit more
adversarial corruption, while strict rules improve estimation accuracy but cause frequent rejections. We propose \mathsf{VISTA}, an adaptive
algorithm that tunes the acceptance threshold using the optimization history. Numerical results show that \mathsf{VISTA} improves convergence over static thresholds. We also provide a rigorous convergence analysis showing that, with suitable incentive-aware adaptation, adversary-dominated decentralized learning can retain the asymptotic convergence behavior of standard SGD without relying on an honest majority.
Read the full abstract ▾
All authors Hanzaleh Akbari Nodehi, Parsa Moradi, Soheil Mohajer, Mohammad Ali Maddah-Ali
Machine Learning
RelAgent: LLM Agents as Data Scientists for Relational Learning
Xingyue Huang, Louis Tichelman, Jinwoo Kim et al. (+2 more)
Relational learning is a challenging problem that has motivated a wide range of approaches, including graph-based models (e.g., graph neural networks, graph
transformers), tabular methods (e.g., tabular foundation models), and sequence-based approaches (e.g., large
language models), each with its own advantages and limitations. We propose RelAgent, an LLM-based autonomous data scientist for relational learning, which operates in two phases. In the search phase, an LLM
agent uses database, validation, and evaluation workspace tools to construct SQL feature programs and select a predictive model. In the inference phase, the resulting program is executed without further LLM calls. The final predictor consists of SQL queries and a classical model, enabling fast, deterministic, and intrinsically interpretable predictions: features are human-readable queries, and predictions depend only on the resulting query-defined feature map, enabling scalable deployment using standard database systems.
Read the full abstract ▾
All authors Xingyue Huang, Louis Tichelman, Jinwoo Kim, Krzysztof Olejniczak, İsmail İlkan Ceylan
Machine Learning · Cross-listed from q-bio.QM
PPI-Net connects molecular protein interactions to functional processes in disease
Kyle Higgins, Guadalupe Gonzalez, Dennis Veselkov et al. (+2 more)
Understanding how molecular alterations propagate across biological systems to drive disease remains a central challenge. Although high-throughput profiling enables comprehensive characterization of tumor states, most models neglect structured biological relationships or lack interpretability across scales. Here we present PPI-Net, a hierarchical graph neural network that integrates
protein-
protein interaction (PPI) networks with pathway-level representations to model disease from molecular interactions to functional processes. Patient-specific molecular profiles are embedded within a shared interaction network from STRING and propagated through a multi-layer Reactome hierarchy using graph
attention, enabling aggregation of gene-level signals into higher-order biological programs. Across RNA-seq data from ten cancer types from The Cancer Genome Atlas, PPI-Net achieves robust predictive performance, with balanced accuracy exceeding 90% in multiple cohorts. Comparative analysis on RNA-Seq data from breast cancer demonstrated that PPI-Net's integration of the Reactome hierarchy improved balanced accuracy by 6.7% relative to a PPI-only model, while hierarchical multi-level supervision improved balanced accuracy by 12.3% relative to using only a single top-level prediction head. Applying a multi-omics approach using RNA-seq and methylation data improves model interpretation, recovering canonical oncogenic modules, including TP53-AKT signaling and stress response pathways, while revealing convergence onto coherent programs such as ion signaling and cellular responses to stimuli. These results demonstrate that integrating interaction networks with pathway hierarchies enables accurate prediction while providing mechanistic insight into cancer biology.
Read the full abstract ▾
All authors Kyle Higgins, Guadalupe Gonzalez, Dennis Veselkov, Ivan Laponogov, Kirill Veselkov
Machine Learning
Approximation-Free Differentiable Oblique Decision Trees
Subrat Prasad Panda, Blaise Genest and Arvind Easwaran
Decision Trees (DTs) are widely used in safety-critical domains such as medical diagnosis, valued for their interpretability and effectiveness on tabular data. However, training accurate oblique DTs is challenging due to complex optimization landscapes and overfitting risks, particularly in regression. Recent advances have introduced differentiable formulations that enable gradient-based training and joint optimization of decision boundaries and leaf regressors. Yet, existing approaches typically rely on approximations, either through probabilistic softening of boundaries (soft DTs) or quantized gradients such as the Straight-Through
Estimator (STE). To overcome these limitations, we propose DTSemNet, a novel, semantically equivalent, and invertible representation of hard oblique DTs as neural networks. DTSemNet enables end-to-end training with standard gradient descent, eliminating the need for approximations in both classification and regression. While classification aligns naturally with this formulation, regression remains challenging due to the joint optimization of internal nodes and leaf regressors. To address this, we analyze the limitations of STE and introduce an annealed Top-k method that provides accurate gradient signals without approximation. Extensive experiments on classification and regression
benchmarks show that DTSemNet-trained oblique DTs outperform state-of-the-art differentiable DTs. Furthermore, we demonstrate that DTSemNet can serve as programmatic DT policies in reinforcement learning environments, thereby broadening their applicability.
Read the full abstract ▾
Machine Learning · Cross-listed from math.NA
NSPOD: Accelerating Krylov solvers via DeepONet-learned POD subspaces
Francesc Levrero-Florencio, Youngkyu Lee, Jay Pathak et al. (+1 more)
The convergence of Krylov-based linear iterative solvers applied to parametric partial differential equations (PDEs) is often highly sensitive to the domain, its discretization, the location/values of the applied Dirichlet/Neumann boundary conditions, body forces and material properties, among others. We have previously introduced hybridization of classical linear iterative solvers with neural operators for specific geometries, but they tend to not perform well on geometries not previously seen during training. We partially addressed this challenge by introducing the deep operator network Geo-DeepONet and hybridizing it with Krylov-based iterative linear solvers, which, despite learning effectively across arbitrary unstructured meshes without requiring retraining, led to only modest reductions in iterations compared to state-of-the-art preconditioners. In this study we introduce Neural Subspace Proper Orthogonal Decomposition (NSPOD), a multigrid-like deep operator network-based preconditioner which can dramatically reduce the number of iterations needed for convergence in Krylov-based linear iterative solvers, even when compared to state-of-the-art methods such as algebraic multigrid preconditioners. We demonstrate its efficiency via numerical experiments on a linearized version of solid mechanics PDEs applied to unstructured domains obtained from complex CAD geometries. We expect that the findings in this study lead to more efficient hybrid preconditioners that can match, or possibly even surpass, the convergence properties of the current gold standard preconditioning methods for solid mechanics PDEs.
Read the full abstract ▾
All authors Francesc Levrero-Florencio, Youngkyu Lee, Jay Pathak, George Em Karniadakis
Machine Learning
Scaling Categorical Flow Maps
Oscar Davis, Anastasiia Filippova, Pierre Ablin et al. (+4 more)
Continuous
diffusion and flow matching models could represent a powerful alternative to autoregressive approaches for
language modelling (LM), as they unlock a host of advantages currently reserved for continuous modalities, including accelerated sampling and tilting. Recently, several works have demonstrated the possibility of generating discrete data continuously by a simple flow matching process between a Gaussian and the one-hot encoded data distribution. They have further shown the feasibility of accelerated sampling via Categorical Flow Maps (CFMs), resulting in competitive sample quality in the few-step regime. However, this method had only been evaluated at relatively modest scales ($<1$B), leaving the question of its scalability completely open. In this article, we train a $1.7$B-parameter base flow model on $2.1$T
tokens and self-distill it into a CFM that generates diverse, high-quality text in as few as $4$ inference steps while maintaining near-data-level
token entropy. Furthermore, we introduce a likelihood bound for CFMs in the semi-discrete setting, and show that they can be used to score the model on standard LM
benchmarks, achieving results in the same range as discrete
diffusion methods. Finally, we uncover some of the challenges that arise from training these models at scale, and we provide prescriptive insights on loss weighting and time scheduling.
Read the full abstract ▾
All authors Oscar Davis, Anastasiia Filippova, Pierre Ablin, Victor Turrisi, Amitis Shidani, Marco Cuturi, Louis Béthune
Machine Learning
OrScale: Orthogonalised Optimization with Layer-Wise Trust-Ratio Scaling
Yuxuan Lou and Yang You
Muon fixes the direction of every
matrix-valued update at the polar factor of its momentum, while each layer's step magnitude is addressed only by a static shape correction. We derive a dynamic per-layer scalar by adapting the LARS/LAMB trust-ratio principle to the orthogonalized setting, where the standard denominator candidates—the raw momentum norm or the polar-factor norm—either live in the wrong unit space or carry no update-scale information. The resulting method, OrScale, uses the norm of the parameter-space direction actually applied and anchors each layer's ratio at one via a per-layer calibration, so that the Moonlight recipe (tuned for AdamW, shared with
Muon via RMS matching) transfers with no additional sweep; a component ablation confirms each design choice is individually load-bearing. Theoretically, OrScale retains a
nuclear-norm $O(1/\sqrt{T})$ convergence rate for any clipped multiplier and achieves a strict layer-adaptive descent gain $\kappa_{\mathrm{eff}}>1$ under two conditions estimable from standard training diagnostics—a bound that predicts the gain should grow with architectural heterogeneity. Experiments confirm the prediction: with every hyperparameter inherited verbatim from the Moonlight recipe, OrScale matches or beats
Muon+Moonlight across dense 125M–1.1B FineWeb-Edu pre-training, and on a 16B-A3B mixture-of-experts model—where the logged trust ratios separate cleanly by layer class—the gap widens by an order of magnitude to $0.130$ nats ($3.8\%$ relative) at parity wall-clock cost.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CR
GRASP – Graph-Based Anomaly Detection Through Self-Supervised Classification
Robin Buchta, Carsten Kleiner, Felix Heine et al. (+1 more)
Advanced persistent threat (APT) attacks remain difficult to detect due to their stealth, adaptability, and use of legitimate system components. Provenance-based intrusion detection systems (PIDS) offer a promising defense by capturing detailed relationships between system components and actions. However, current PIDS rely on predefined or subset-determined thresholds, which limit detection stability and the ability to detect any anomalous behavior in general. Furthermore, related work often neglects the role of process executables, which describe system activity by interacting through a process with files, network components, and other processes. We introduce GRASP, a PIDS based on masked self-supervised classification. GRASP masks the executable information of processes and learns to infer it from their two-hop provenance graph neighborhood, marking misclassified processes as anomalies. It captures behavior patterns for the learned executables without thresholding, making it robust against interference and unknown activities. Evaluations on the DARPA TC and OpTC
datasets demonstrate that GRASP consistently detects anomalous behavior, including known attack-related activities, outperforming existing systems. Our PIDS identifies all documented attacks on
datasets where the behavior of executables is learnable. In addition, compared to existing systems, GRASP uncovers potentially malicious anomalous behavior not labeled as an attack in the documentation.
Read the full abstract ▾
All authors Robin Buchta, Carsten Kleiner, Felix Heine, Gabi Dreo Rodosek
Machine Learning
The Minimax Rate of Perturbed Second-Order Calibration
Kamil Ciosek, Banafsheh Rafiee, Sina Ghiassian et al. (+1 more)
Second-order calibration error quantifies how closely a higher-order predictor's epistemic-uncertainty estimate matches the conditional variance of the label probability on its level sets. We characterize the minimax rate of estimating the second-order calibration error for binary classification in the regime where a small perturbation is applied to the classifier outputs. Our procedure is simple: add independent bandwidth-$h$ sech noise to the score coordinates, then regress $Y^{(1)}$ and $Y^{(1)}Y^{(2)}$ on the perturbed score using low-degree
polynomials. Crucially, the sech perturbation makes the calibration functions analytic in a suitable strip. The resulting
estimator has error $O_h(\log^{3/2}n/\sqrt n)$, with explicit constants. In the same setting, a matching $\Omega(1/\sqrt{n})$ lower bound establishes minimax optimality up to logarithmic factors. As a corollary, we give a finite-sample guarantee for perturbed second-order Platt scaling, yielding a post-hoc procedure that recalibrates both the mean prediction and the epistemic-variance estimate of the perturbed higher-order predictor. Along the way, we give an explicit two-moment formulation of second-order calibration error and relate it quantitatively to the bucketed formulation of Ahdritz et al. [2025]. Our experiments confirm the predicted rate and the quality of the recalibrated uncertainties.
Read the full abstract ▾
All authors Kamil Ciosek, Banafsheh Rafiee, Sina Ghiassian, Nicolò Felicioni
Machine Learning · Cross-listed from cs.CV
Text-to-CAD Evaluation with CADTests
Dimitrios Mallis, Marco Wang, Ahmet Serdar Karadeniz et al. (+3 more)
Text-to-CAD has recently emerged as an important task with the potential to substantially accelerate design workflows. Despite its significance, there has been surprisingly little work on Text-to-CAD evaluation, and assessing CAD model generation performance remains a considerable challenge. In this work, we introduce a new evaluation perspective for Text-to-CAD based on automated testing. We propose CADTestBench, the first test-based
benchmark for Text-to-CAD, based on CADTests, executable software tests that verify whether a generated CAD model satisfies the
geometric and topological requirements of the input
prompt. Using CADTestBench, we conduct comprehensive benchmarking of recent Text-to-CAD methods and further demonstrate that CADTests can also guide CAD model generation, yielding simple baselines that surpass performance of current methods. CADTestBench code and data are available at GitHub and Hugging Face
dataset.
Read the full abstract ▾
All authors Dimitrios Mallis, Marco Wang, Ahmet Serdar Karadeniz, Elisa Ricci, Anis Kacem, Djamila Aouada
Machine Learning · Cross-listed from cs.CL
Beyond Confidence: Rethinking Self-Assessments for Performance Prediction in LLMs
Sree Bhattacharyya, Samarth Khanna, Leona Chen et al. (+3 more)
Large
Language Models (LLMs) are increasingly used in settings where reliable self-assessment is critical. Assessing model reliability has evolved from using probabilistic correctness estimates to, more recently, eliciting verbalized confidence. Confidence, however, has been shown to be an inconsistent and overoptimistic predictor of model correctness. Drawing on cognitive appraisal theory, a framework from human psychology that decomposes self-evaluation into multiple components, we propose a multidimensional perspective on model self-assessment. We elicit six appraisal-based dimensions of self-assessment, alongside confidence, and evaluate their utility for predicting model failure across 12 LLMs and 38 tasks spanning eight domains. We find that competence-related appraisal dimensions, particularly effort and ability, consistently match or outperform confidence across most settings. Effort additionally yields less overoptimistic estimates that remain stable across model sizes. In contrast, affective dimensions provide marginally predictive signals. Furthermore, the most informative dimension varies systematically with task characteristics: effort is most predictive for
reasoning-intensive tasks, while ability and confidence dominate on
retrieval-oriented tasks. Broadly, our findings indicate that structured multidimensional self-assessment is a promising approach to improving the reliability and safety of
language model deployment across diverse real-world settings.
Read the full abstract ▾
All authors Sree Bhattacharyya, Samarth Khanna, Leona Chen, Lucas Craig, Tharun Dilliraj, James Z. Wang
Machine Learning
Flexible Routing via Uncertainty Decomposition
Charlotte Peale, Siddartha Devic, Parikshit Gopalan et al. (+2 more)
A key strategy for balancing performance and cost in modern machine learning systems is to dynamically route queries to either a low-cost model or a more expensive oracle (such as a large pretrained model or human expert), an approach known as model routing. In this work we present a new uncertainty-aware router that (1) avoids unnecessary oracle calls on inherently ambiguous queries, and (2) adapts dynamically to different loss functions and cost parameters through simple hyperparameter changes, without retraining. Our method, applicable to any classification setting where multiple independent annotations per input are available, is based on decomposing total uncertainty into irreducible and reducible components using higher-order predictors [Ahdritz et al., 2025]. This enables a unified approach to both routing and abstention: predict with the weak model when uncertainty is low, route to the oracle when reducible uncertainty is high, and abstain when irreducible uncertainty is high. Our router comes with strong theoretical guarantees bounding regret relative to optimal task-specific routers. We conduct experiments on both synthetic and real-world
datasets that demonstrate the benefits of our approach in suitable regimes – in particular, whenever reducible and irreducible uncertainty are not too correlated.
Read the full abstract ▾
All authors Charlotte Peale, Siddartha Devic, Parikshit Gopalan, Udi Wieder, Aravind Gollakota
Machine Learning
Prune-OPD: Efficient and Reliable On-Policy Distillation for Long-Horizon Reasoning
Zhicheng Yang, Zhijiang Guo, Yifan Song et al. (+5 more)
On-policy
distillation (OPD) leverages dense teacher rewards to enhance
reasoning models. However, scaling OPD to long-horizon tasks exposes a critical flaw: as the student's generated prefix inevitably diverges from the teacher's thought process, the teacher's dense reward loses local exploitability. Continuing to generate and evaluate
tokens on these “drifted” trajectories not only degrades reward quality but also incurs massive computational waste. To address this, we introduce Prune-OPD, a framework that dynamically aligns training budgets with supervision quality. By continuously monitoring the local compatibility between student and teacher predictions (e.g., via top-$k$ overlap), Prune-OPD detects prefix-drift events in real time. Upon detecting severe drift, it monotonically down-weights subsequent unreliable rewards and triggers dynamic rollout truncation. This allows the training process to halt futile generation and reallocate compute strictly to reliable teacher supervision. Across diverse teacher-student combinations, Prune-OPD consistently aligns computation with supervision reliability. When prefix drift makes dense teacher rewards unreliable, it reduces training time by 37.6%–68.0% while preserving, and often improving, performance on challenging
benchmarks (AMC, AIME, HMMT). When student-teacher compatibility remains high, it automatically preserves long-context supervision by expanding the training window. These results suggest that Prune-OPD improves OPD not by blindly shortening rollouts, but by reallocating computation toward locally exploitable teacher rewards.
Read the full abstract ▾
All authors Zhicheng Yang, Zhijiang Guo, Yifan Song, Minrui Xu, Yongxin Wang, Yiwei Wang, Xiaodan Liang, Jing Tang
Machine Learning
Toward Privileged Foundation Models:LUPI for Accelerated and Improved Learning
Xueying Ding and Leman Akoglu
Training foundation models is computationally intensive and often slow to converge. We introduce PIQL,Privileged Information for Quick and Quality Learning, the first framework to systematically integrate privileged information (PI) to simultaneously accelerate learning and improve generalization in tabular foundation models (TFMs). We construct two complementary forms of PI: (i) aggregate
dataset-level statistics that reduce the burden on in-context learning, and (ii) encodings of the underlying data-generating program, providing knowledge beyond observable data. We further design an architecture that effectively transfers the train-time-only PI by learning to reconstruct it from observed context at inference. We provide a theoretical analysis characterizing conditions under which PI reduces the population-level approximation gap and accelerates convergence in finite-data regimes. Empirical evidence shows that PIQL enables TFMs to achieve faster convergence, lower final loss, and better generalization, in effect, reducing data and compute requirements. Our work establishes PI-guided pretraining as a principled and practical paradigm for improving the efficiency and performance of foundation models.
Read the full abstract ▾
Machine Learning
Neural Operators as Efficient Function Interpolators
Vasilis Niarchos, Angelos Sirbu and Sokratis Trifinopoulos
Neural operators (NOs) are designed to learn maps between infinite-dimensional function spaces. We propose a novel reframing of their use. By introducing an auxiliary base-space, any finite-dimensional function can be viewed as an operator acting by composition on functions of the base-space. Through a range of
benchmarks on analytic functions of increasing complexity and dimensionality, we demonstrate that NOs can match or outperform standard multilayer perceptrons and Kolmogorov–Arnold Networks in accuracy while requiring significantly fewer parameters and training time. As a real-world application, we apply a two-dimensional Tensorized
Fourier Neural Operator (TFNO) to the
nuclear chart, learning a correction to state-of-the-art
nuclear mass models as a partially observed residual field. A TFNO ensemble reaches a held-out root-mean-square error of 198.2 keV, placing it among the best recent neural-network approaches while retaining high parameter efficiency and short training times. More broadly, these results introduce NOs as a scalable framework for finite-dimensional function interpolation, from analytic
benchmarks to structured scientific data.
Read the full abstract ▾
Machine Learning
Hessian Surgery: Class-Targeted Post-Hoc Rebalancing via Hessian Spike Perturbation
Hugo Vigna and Samuel Bontemps
The Hessian spectrum of trained deep networks exhibits a characteristic structure: a continuous bulk of near-zero
eigenvalues and a small number of large outlier
eigenvalues (spikes), confirming the relevance of Random
Matrix Theory in deep learning. The spike count matches the number of classes minus one. While prior work has described this structure, no method has exploited it operationally to improve classification performance. We propose Hessian Surgery, a post-hoc optimization method that directly perturbs model weights along spike eigenvectors to rebalance per-class accuracy without retraining. We introduce (i) a spike-class sensitivity
matrix that quantifies the directional derivative of each class's accuracy along each spike eigenvector, (ii) a constrained optimization of perturbation coefficients that targets weak classes while preserving strong ones, and (iii) an adaptive amplitude control that raises or lowers the perturbation budget based on iteration-level improvement signals. We obtain encouraging results on CIFAR-10 and ISIC-2019 on both balanced accuracy and standard deviation.
Read the full abstract ▾
Machine Learning
Nils Grünefeld, Bertram Højer, Philipp Mondorf et al. (+5 more)
Language model (LM) "
reasoning", commonly described as Chain-of-Thought or test-time scaling, often improves
benchmark performance, but the dynamics underlying this process remain poorly understood. We study these dynamics through the lens of uncertainty quantification by treating the "
reasoning" traces, the intermediate
token sequences generated by LMs, as evolving model states. We summarize each trace by an uncertainty trace profile: a small set of features describing the shape of the uncertainty signal over its trace, such as its slope and linearity. We find that across five LMs evaluated on GSM8K and ProntoQA, these profiles predict whether a trace yields a correct final answer with AUROC up to 0.807, improving markedly on recent related work. We reach AUROC 0.801 using only the first few hundred
tokens of full traces, suggesting that errors can be detected early in the generation. A detailed comparison of correct and incorrect traces further reveals qualitatively distinct uncertainty profiles, with correct traces showing a steeper and less linear decline in uncertainty. Together, the results suggest that our method, grounded in decision-making under uncertainty, provides a principled lens for studying the generative process underlying LM "
reasoning".
Read the full abstract ▾
All authors Nils Grünefeld, Bertram Højer, Philipp Mondorf, Barbara Plank, Anna Rogers, Christian Hardmeier, Stefan Heinrich, Jes Frellsen
Machine Learning
POETS: Uncertainty-Aware LLM Optimization via Compute-Efficient Policy Ensembles
Nicolas Menet, Andreas Krause and Abbas Rahimi
Balancing exploration with exploitation is a core challenge in black-box optimization and reinforcement learning. We introduce POETS ($\textbf{Po}$licy $\textbf{E}$nsembles for $\textbf{T}$hompson $\textbf{S}$ampling), a novel framework that bridges uncertainty quantification and policy optimization. Our approach is grounded in the insight that policies trained with Kullback-Leibler (KL) regularization implicitly encode an underlying reward function. Building on this, POETS bypasses the complex, nested process of training an uncertainty-aware reward model and separately fitting a policy to this model. Instead, we directly train a policy ensemble to capture epistemic uncertainty by matching implicitly encoded reward functions to online, bootstrapped data. To overcome the prohibitive compute and
memory constraints of ensembling Large
Language Models (LLMs), POETS utilizes an efficient architecture: the ensemble shares a pre-trained backbone while maintaining diversity through independent Low-Rank Adaptation (LoRA) branches. Theoretically, we show that POETS implicitly approximates KL-regularized Thompson sampling, a principled optimization strategy with optimal cumulative regret bounds of $O(\sqrt{T \gamma_T})$. Empirically, we demonstrate that POETS achieves state-of-the-art sample efficiency across diverse scientific discovery domains, including
protein search, therapeutic molecular design, and
quantum circuit design. POETS also improves the optimization trajectory and out-of-distribution transfer of reinforcement learning across mathematical
reasoning and code generation tasks, proving particularly robust in off-policy settings with experience replay or in small
dataset regimes.
Read the full abstract ▾
Machine Learning
Training-Induced Escape from Token Clustering in a Mean-Field Formulation of Transformers
Noboru Isobe, Daisuke Inoue and Masaaki Imaizumi
Transformers perform inference by iteratively transforming
token representations across layers. This layerwise computation has been studied empirically, and recent mean-field theories of
Transformer dynamics explain how
attention can drive
token distributions toward clustering. However, existing analyses treat model parameters as prescribed, leaving open how training reshapes this clustering picture. We study this question in a noisy mean-field
Transformer in which only a parameter-linear FFN is trained under $L^2$ regularization. We identify a training-induced phase of the inference dynamics across layers: after clustering in the early layers, the
token distribution can leave the clustered regime near the final layers. A single gradient-descent update already produces this escape, and minimizers of the training objective satisfy an exponential turnpike bound under which clustering persists through the interior layers. Our analysis is based on an
entropy-regularized interaction energy that captures the clustering bias of
attention, for which we establish a double-well landscape and a local Polyak-Łojasiewicz inequality. More broadly, our results point toward a training-aware mean-field theory in which training and inference dynamics of
Transformers are treated together.
Read the full abstract ▾
Machine Learning · Cross-listed from eess.SY
Interactive Trajectory Planning with Learning-based Distributionally Robust Model Predictive Control and Markov Systems
Erik Börve, Nikolce Murgovski, Morteza Haghir Chehreghani et al. (+1 more)
We investigate interactive trajectory planning subject to uncertainty in the decisions of surrounding
agents. To control the ego-
agent, we aim to first learn the decision distribution and solve a
Stochastic Model Predictive Control (SMPC) problem. To account for errors in the learned distribution, we show that it is possible to utilize Probably Approximately Correct (PAC) learning in combination with Distributionally Robust (DR) optimization to obtain a solution which accounts for the errors induced by the learning model. The results indicate that our PAC learning-based DR-MPC framework provides a method to interpolate between a robust MPC and an omnipotent SMPC, based on the available number of samples.
Read the full abstract ▾
All authors Erik Börve, Nikolce Murgovski, Morteza Haghir Chehreghani, Leo Laine
Machine Learning
Pre-trained Tabular Foundation Models as Versatile Summary Networks for Neural Posterior Estimation
Elliot Pickens, Chiraag Gohel and Sidharth Satya
In this work, we study TabPFN as a training-free, modular summary network for simulation-based
Bayesian inference (SBI). Tabular foundation models such as TabPFN are pretrained on broad families of synthetic tabular data-generating processes and adapt at test time through in-context learning, making them natural candidates for SBI, where posterior estimation often depends on learning informative summaries of simulated observations. We propose PFN-NPE: a general recipe that uses a pretrained TabPFN encoder as a fixed summary network for simulator outputs, then pairs the resulting summaries with a downstream inference head chosen for the problem. With normalizing flows as the default inference head, PFN-NPE matches established posterior approximation methods and sometimes outperforms them. More importantly, diagnostic probes show that the TabPFN-derived summaries often preserve useful posterior location and marginal information. These analyses also reveal a limitation in that TabPFN-derived summaries may struggle to represent the joint posterior structure even when the marginals are well recovered. Still, our experiments show that TabPFN can serve as an effective summary network across a diverse set of SBI settings, with the inference network left modular and task-dependent.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.FL
SMT-Based Active Learning of Weighted Automata
Tiago Ferreira, Kevin Batz and Alexandra Silva
We present an SMT-based active learning
algorithm for nondeterministic weighted automata (WFAs) as a practical and robust alternative to Hankel/L*-style methods. Our
algorithm is parametric in a given semiring and, if it terminates, guaranteed to produce minimal WFAs. We prove partial correctness and provide a sufficient termination condition, which in particular implies termination for all finite semirings. Our extensive experimental evaluation shows that our
algorithm is capable of learning numerous minimal WFAs over both finite and infinite semirings, vastly outperforms a naive baseline, and is competitive with a state-of-the-art
algorithm while producing significantly smaller automata and requiring less interaction with the teacher.
Read the full abstract ▾
Machine Learning
Efficient Verification of Neural Control Barrier Functions with Smooth Nonlinear Activations
Jun Zhang, Haibo Zhang, Chun Liu et al. (+2 more)
Formal verification of neural control barrier functions (NCBFs) remains challenging, especially for neural networks with nonlinear activations like \(\tanh\). Existing CROWN-based methods rely on conservative linear relaxations for Jacobian bounds, limiting scalability. We propose LightCROWN, which computes tighter Jacobian bounds by exploiting the analytical properties of activation functions. Experiments on nonlinear control systems including the inverted pendulum, Dubins car, and planar quadrotor demonstrate that LightCROWN improves verification success rates up to 100%, while enhancing speed and scalability. Our approach provides a generalizable improvement for CROWN-based frameworks, enabling more efficient verification of complex NCBFs. The code can be found at github.com/Autonomous-Systems-and-Control-Lab/verify-neural-CBF.
Read the full abstract ▾
All authors Jun Zhang, Haibo Zhang, Chun Liu, Xiaofan Wang, Liang Xu
Machine Learning
When Losses Align: Gradient-Based Composite Loss Weighting for Efficient Pretraining
Ivan Karpukhin and Andrey Savchenko
Modern deep models are often pretrained on large-scale data with missing labels using composite objectives, where the relative weights of multiple loss terms act as hyperparameters. Tuning these weights with random search or
Bayesian optimization is computationally expensive, as it requires many independent training runs. To address this, we propose a gradient-based bilevel method that learns pretraining loss weights online by aligning the composite pretraining gradient with a downstream objective. By exploiting the structure of the loss, the method avoids the multiple backward passes typically required by truncated backpropagation through the full model, reducing the overhead of hyperparameter tuning to approximately 30% above a single training run. We evaluate the approach on event-sequence modeling and self-supervised computer
vision, where it matches or improves upon carefully tuned baselines while substantially reducing the cost of hyperparameter tuning compared to random or
Bayesian search.
Read the full abstract ▾