Machine Learning
Unifying Goal-Conditioned RL and Unsupervised Skill Learning via Control-Maximization
Alireza Modirshanechi, Benjamin Eysenbach, Peter Dayan et al. (+1 more)
Unsupervised pretraining has driven empirical advances in goal-conditioned reinforcement learning (GCRL), but its theoretical foundations remain poorly understood. In particular, an influential class of methods, mutual information skill learning (MISL), discovers behaviorally diverse skills that can later be used for downstream goal-reaching. However, it remains a theoretical mystery why skills learned through MISL should support goal-reaching. A subtle challenge is that both GCRL and MISL are umbrella terms: Different GCRL tasks use distinct criteria for measuring goal-reaching performance, while different MISL methods optimize distinct notions of behavioral diversity. We address this challenge and unify GCRL and MISL as instances of control maximization. We identify three canonical GCRL formulations and prove that they are fundamentally inequivalent: They can induce incompatible optimal policies even in the same environment. Nevertheless, they all share a common interpretation: A well-performing goal-conditioned policy is one whose future trajectory is highly sensitive to the commanded goal, with the precise notion of sensitivity determined by the GCRL formulation. Noting that MISL objectives can be understood as measures of skill sensitivity akin to goal sensitivity, we show that MISL objectives are bounded by formulation-specific downstream goal sensitivities. These bounds establish a precise correspondence between MISL methods and downstream GCRL tasks: For every GCRL formulation, there exists a matching MISL objective for which more diverse skills afford greater downstream goal sensitivity. Our results thus lay a theoretical foundation for RL pretraining and have important practical implications, such as suggesting which pretraining objectives to use when a user cares about a specific class of downstream tasks.
Read the full abstract ▾
All authors Alireza Modirshanechi, Benjamin Eysenbach, Peter Dayan, Eric Schulz
Machine Learning
Matrix-Valued Optimism is Matrix-Valued Augmentation: Additive Hybrid Designs for Constrained Optimization
Jiayi Zhao
Augmented Lagrangian and optimistic primal–dual methods stabilize equality-constrained optimization through seemingly different mechanisms: the former adds constraint-dependent primal
curvature, while the latter adds dual
memory. Recent work has shown that these mechanisms are equivalent for scalar parameters. We extend this equivalence to
matrix-valued correction. We prove an additivity principle: for symmetric
matrix parameters, the ideal primal trajectory depends only on the summed correction
matrix, not on how it is split between augmented and optimistic channels. This exposes a design freedom: algebraically equivalent decompositions can have different finite-step feasibility because augmented correction affects primal
curvature, whereas optimistic correction affects the scale of the dual
memory correction. We formulate the resulting step-size-limited design problem and derive a closed-form hybrid rule that selects a
matrix correction, splits it between the two channels, and chooses primal and dual steps using local spectral weights. Experiments on nonlinear equality-constrained problems with controlled constraint-Jacobian conditioning show that the hybrid design improves over pure augmented and pure optimistic endpoints, closely tracks a grid-search hybrid oracle, and is competitive with first-order primal–dual baselines under mild-to-moderate ill-conditioning. The experiments also identify the expected limitation: exact cancellation requires increasingly large
matrix corrections as the constraint Jacobian becomes ill-conditioned.
Read the full abstract ▾
Machine Learning
SymDrift: One-Shot Generative Modeling under Symmetries
Samir Darouich, Vinh Tong, Lluís Pastor-Pérez et al. (+3 more)
Generative modeling of physical systems, such as molecules, requires learning distributions that are invariant under global
symmetries, such as rotations in three-dimensional space. Equivariant
diffusion and flow matching models can incorporate such invariances effectively, even when trained on a non-invariant empirical distribution, but they typically rely on costly multi-step sampling. Recently, drifting models have emerged as an efficient alternative, enabling single-step generation and achieving state-of-the-art performance in generative modeling tasks. However, we show that drifting models face a
symmetry-specific challenge, since an equivariant generator does not generally produce the same drifting field as the one obtained from the symmetrized target distribution. Addressing this issue would require expensive symmetrization of the empirical distribution. To avoid this cost, we propose SymDrift, a framework that makes the drifting field itself
symmetry-aware. We introduce two complementary strategies: (i) a symmetrized drift in coordinate space based on optimal
alignment, and (ii) a $G$-invariant
embedding that removes
symmetry ambiguity by construction. Empirically, SymDrift outperforms existing one-shot methods on standard
benchmarks for conformer and transition state generation, while remaining competitive with significantly more expensive multi-step approaches. By enabling one-shot inference, SymDrift reduces computational overhead by up to 40$\times$ compared to existing baselines, making it promising for high-throughput applications such as virtual drug screening and large-scale reaction network exploration.
Read the full abstract ▾
All authors Samir Darouich, Vinh Tong, Lluís Pastor-Pérez, Tanja Bien, Loay Mualem, Mathias Niepert
Machine Learning
Listwise Policy Optimization: Group-based RLVR as Target-Projection on the LLM Response Simplex
Yun Qu, Qi Wang, Yixiu Mao et al. (+11 more)
Reinforcement learning with verifiable rewards (RLVR) has become a standard approach for large
language models (LLMs) post-training to incentivize
reasoning capacity. Among existing recipes, group-based policy gradient is prevalent, which samples a group of responses per
prompt and updates the policy via group-relative advantage signals. This work reveals that these optimization strategies share a common
geometric structure: each implicitly defines a target distribution on the response simplex and projects toward it via first-order approximation. Building on this insight, we propose Listwise Policy Optimization (LPO) to explicitly conduct the target-projection, which demystifies the implicit target by restricting the proximal RL objective to the response simplex, and then projects the policy via exact divergence minimization. This framework provides (i) monotonic improvement on the listwise objective with bounded, zero-sum, and self-correcting projection gradients, and (ii) flexibility in divergence selection with distinct structural properties through the decoupled projection step. On diverse
reasoning tasks and LLM backbones, LPO consistently improves training performance over typical policy gradient baselines under matched targets, while intrinsically preserving optimization stability and response diversity.
Read the full abstract ▾
All authors Yun Qu, Qi Wang, Yixiu Mao, Heming Zou, Yuhang Jiang, Yingyue Li, Wutong Xu, Lizhou Cai, Weijie Liu, Clive Bai, Kai Yang, Yangkun Chen, Saiyong Yang, Xiangyang Ji
Machine Learning · Cross-listed from cs.CV
Autoregressive Visual Generation Needs a Prologue
Bowen Zheng, Weijian Luo, Guang Yang et al. (+2 more)
In this work, we propose Prologue, an approach to bridging the reconstruction-generation gap in autoregressive (AR) image generation. Instead of modifying visual
tokens to satisfy both reconstruction and generation, Prologue generates a small set of prologue
tokens prepended to the visual
token sequence. These prologue
tokens are trained exclusively with the AR cross-
entropy (CE) loss, while visual
tokens remain dedicated to reconstruction. This decoupled design lets us optimize generation through the AR model's true distribution without affecting reconstruction quality, which we further formalize from an ELBO perspective. On ImageNet 256x256, Prologue-Base reduces gFID from 21.01 to 10.75 without classifier-free guidance while keeping reconstruction almost unchanged; Prologue-Large reaches a competitive rFID of 0.99 and gFID of 1.46 using a standard AR model without auxiliary semantic supervision. Interestingly, driven only by AR gradients, prologue
tokens exhibit emergent semantic structure: linear probing on 16 prologue
tokens reaches 35.88% Top-1, far above the 23.71% of the first 16
tokens from a standard tokenizer; resampling with fixed prologue
tokens preserves a similar high-level semantic layout. Our results suggest a new direction: generation quality can be improved by introducing a separate learned generative representation while leaving the original representation intact.
Read the full abstract ▾
All authors Bowen Zheng, Weijian Luo, Guang Yang, Colin Zhang, Tianyang Hu
Machine Learning · Cross-listed from hep-lat
Diffusion model for SU(N) gauge theories
Javad Komijani, Marina K. Marinkovic and Lara Turgut
Implicit score matching provides a computationally efficient approach for training
diffusion models and generating high-quality samples from complex distributions. In this work, we develop a score-matching framework for SU(N)
lattice gauge theories, which can be extended to other Lie groups. We apply the method to SU(3) gauge configurations with the Wilson gauge action in two and four dimensions and assess the quality of the generated samples by comparison with Hybrid Monte Carlo (HMC) simulations. We show that the
diffusion models can be successfully trained and applied for sampling the Wilson gauge action. For large values of inverse coupling, accurate reverse-time integration requires predictor-corrector schemes, for which we introduce a corrector based on
Hamiltonian molecular dynamics. While the corrector significantly improves sampling quality, it also increases the computational cost. We outline several strategies for improving sampling efficiency.
Read the full abstract ▾
Machine Learning
BoostLLM: Boosting-inspired LLM Fine-tuning for Few-shot Tabular Classification
Yi-Siang Wang, Kuan-Yu Chen, Yu-Chen Den et al. (+1 more)
Large
language models (LLMs) have recently been adapted to tabular prediction by serializing structured features into natural
language, but their performance in low-data regimes remains limited compared to gradient-boosted decision trees (GBDTs). In this work, we revisit the boosting paradigm, traditionally associated with tree ensembles, and ask whether it can be applied as a general training principle for LLM fine-tuning. We propose BoostLLM, a framework that transforms parameter-efficient fine-tuning into a multi-round residual optimization process by training sequential PEFT adapters as weak learners. To incorporate tabular inductive bias, BoostLLM integrates decision-tree paths as a second input view alongside raw features; analysis reveals that the path view acts as a structured teacher in early training steps before the model shifts toward feature-driven representations. Empirically, BoostLLM achieves consistent improvements over standard fine-tuning across multiple LLM backbones and
datasets, matching or surpassing XGBoost across a wide range of shot counts and outperforming GPT-4o-based methods with a 4B model. We further show that the framework scales: pairing with stronger tree models and extended boosting horizons yields additional gains under appropriate stabilization. These results suggest that boosting can serve as a general training principle for LLM fine-tuning, particularly in low-data regimes for structured data.
Read the full abstract ▾
All authors Yi-Siang Wang, Kuan-Yu Chen, Yu-Chen Den, Darby Tien-Hao Chang
Machine Learning
Beyond Autoregressive RTG: Conditioning via Injection Outside Sequential Modeling in Decision Transformer
Yongyi Wang, Hanyu Liu, Lingfeng Li et al. (+6 more)
Decision
Transformer (DT) formulates offline reinforcement learning as autoregressive sequence modeling, achieving promising results by predicting actions from a sequence of Return-to-Go (RTG), state, and action
tokens. However, RTG is a scalar that summarizes future rewards, containing far less information than typical state or action vectors, yet it consumes the same computational budget per
token. Worse, the self-
attention cost of
Transformers grows quadratically with sequence length, so including RTG as a separate
token adds unnecessary overhead. We propose SlimDT, which removes RTG from the autoregressive sequence. Instead, we inject RTG information into the state representations before the sequential modeling step, allowing the
Transformer to process only a compact (state, action) sequence. This reduces the sequence length by one-third, directly improving inference efficiency. On the D4RL
benchmark, SlimDT surpasses standard DT across various tasks and achieves performance comparable to existing state-of-the-art methods. Decoupling a sparse conditioning signal from an information-rich sequence thus yields both computational gains and higher task performance.
Read the full abstract ▾
All authors Yongyi Wang, Hanyu Liu, Lingfeng Li, Bozhou Chen, Ang Li, Qirui Zheng, Xionghui Yang, Chucai Wang, Wenxin Li
Machine Learning · Cross-listed from eess.SP
CredibleDFGO: Differentiable Factor Graph Optimization with Credibility Supervision
Liang Qian, Penggao Yan, Penghui Xu et al. (+1 more)
Global
navigation satellite system (GNSS) positioning is widely used for urban
navigation, but the covariance reported by the GNSS solver is often unreliable in urban canyons. Existing differentiable factor graph optimization (DFGO) methods learn measurement weighting through the solver, but they still use position-only objectives. As a result, the position estimate may improve while the reported covariance remains too small, too large, or incorrectly oriented. We propose CredibleDFGO (CDFGO), a differentiable GNSS factor graph framework that makes covariance credibility an explicit training target. A Weighting Generation Network (WGN) predicts per-satellite reliability weights, and a differentiable Gauss-Newton solver maps these weights to a position estimate and a Hessian-derived posterior covariance. We use proper scoring rules to supervise the East-North predictive distribution end to end. We study negative log-likelihood (NLL), the energy score (ES), and their combination. Results on three UrbanNav test scenes show consistent gains in covariance credibility. Positioning accuracy also improves on the medium-urban and harsh-urban scenes; on the deep-urban scene, both the mean horizontal error and the 95th-percentile error improve. On the harsh-urban Mong Kok (MK) scene, CDFGO-Combined reduces the mean horizontal error from 13.77 m to 11.68 m, reduces NLL from 40.63 to 6.59, and reduces ES from 12.31 to 9.05 relative to DFGO (MAE). Case studies link the MK improvement to better axis-wise consistency, more credible local covariance ellipses, and satellite-level reweighting.
Read the full abstract ▾
All authors Liang Qian, Penggao Yan, Penghui Xu, Li-Ta Hsu
Machine Learning · Cross-listed from math.ST
Time-Inhomogeneous Preconditioned Langevin Dynamics
Alexander Falk, Laurenz Nagler, Andreas Habring et al. (+1 more)
Langevin sampling from distributions of the form $p(x) \propto \exp(-\Psi(x))$ faces two major challenges: (global) mode coverage and (local) mode exploration. The first challenge is particularly relevant for multi-modal distributions with disjoint modes, whereas the second arises when the potential $\Psi$ exhibits diverse and ill-conditioned local mode geometry. To address these challenges, a common approach is to precondition Langevin dynamics with problem-specific information, such as the sample covariance or the local
curvature of $\Psi$. However, existing preconditioner choices inherently involve a trade-off between global mode coverage and local mode exploration, and no prior method resolves both simultaneously. To overcome this limitation, we propose the TIPreL, which introduces a time- and position-dependent preconditioner. This design effectively addresses both challenges mentioned above within a single framework. We establish convergence of the resulting dynamics in the Wasserstein-2 distance both in continuous time and for a tamed Euler discretization. In particular, our analysis extends the existing state of the art by proving convergence under time- and space-dependent
diffusion coefficients, and only locally Lipschitz drifts, which has not been covered by prior work. Finally, we experimentally compare TIPreL with competing preconditioning schemes on a two-dimensional, severely ill-posed example and on a
Bayesian logistic regression task in higher dimensions, confirming the efficiency of the proposed method.
Read the full abstract ▾
All authors Alexander Falk, Laurenz Nagler, Andreas Habring, Thomas Pock
Machine Learning · Cross-listed from cs.CV
Revisiting Uncertainty: On Evidential Learning for Partially Relevant Video Retrieval
Jun Li, Peifeng Lai, Xuhang Lou et al. (+5 more)
Partially relevant
video retrieval aims to retrieve untrimmed
videos using text queries that describe only partial content. However, the inherent asymmetry between brief queries and rich
video content inevitably introduces uncertainty into the
retrieval process. In this setting, vague queries often induce semantic ambiguity across
videos, a challenge that is further exacerbated by the sparse temporal supervision within
videos, which fails to provide sufficient matching evidence. To address this, we propose Holmes, a hierarchical evidential learning framework that aggregates multi-granular cross-modal evidence to quantify and model uncertainty explicitly. At the inter-
video level, similarity scores are interpreted as evidential support and modeled via a Dirichlet distribution. Based on the proposed three-fold principle, we perform fine-grained query identification, which then guides query-adaptive calibrated learning. At the intra-
video level, to accumulate denser evidence, we formulate a soft query-clip
alignment via flexible optimal transport with an adaptive dustbin, which alleviates sparse temporal supervision while suppressing spurious local responses. Extensive experiments demonstrate that Holmes outperforms state-of-the-art methods. Code is released at https://github.com/lijun2005/ICML26-Holmes.
Read the full abstract ▾
All authors Jun Li, Peifeng Lai, Xuhang Lou, Jinpeng Wang, Yuting Wang, Ke Chen, Yaowei Wang, Shu-Tao Xia
Machine Learning · Cross-listed from cs.AR
PoTAcc: A Pipeline for End-to-End Acceleration of Power-of-Two Quantized DNNs
Rappy Saha, Jude Haris, Nicolas Bohm Agostini et al. (+2 more)
Power-of-two (PoT) quantization significantly reduces the size of deep neural networks (DNNs) and replaces multiplications with bit-shift operations for inference. Prior work has shown that PoT-quantized DNNs can preserve accuracy for tasks such as image classification; however, their performance on resource-constrained edge devices remains insufficiently understood. While general-purpose edge CPUs and GPUs do not provide optimized backends for bit-shift operations, custom hardware accelerators can better exploit PoT quantization by implementing dedicated shift-based processing elements. However, deploying PoT-quantized models on such accelerators is challenging due to limited support in existing inference frameworks. In addition, the impact of different PoT quantization strategies on hardware design, performance, and energy efficiency during full inference has not been systematically explored. To address these challenges, we propose PoTAcc, an open-source end-to-end pipeline for accelerating and evaluating PoT-quantized DNNs on resource-constrained edge devices. PoTAcc enables seamless preparation and deployment of PoT-quantized models via TensorFlow Lite (TFLite) across heterogeneous platforms, including CPU-only systems and hybrid CPU-FPGA systems with custom accelerators. We design shift-based processing element (shift-PE) accelerators for three PoT quantization methods and implement them on two FPGA platforms. We evaluate accuracy, performance, energy efficiency, and resource utilization across a range of models, including CNNs and
Transformer-based architectures. Results show that our CPU-accelerator design achieves up to 3.6x speedup and 78% energy reduction compared to CPU-only execution for PoT-quantized DNNs on PYNQ-Z2 and Kria boards. The code will be publicly released at https://github.com/gicLAB/PoTAcc
Read the full abstract ▾
All authors Rappy Saha, Jude Haris, Nicolas Bohm Agostini, David Kaeli, José Cano
Machine Learning
Fast Gauss-Newton for Multiclass Cross-Entropy
Mikalai Korbit and Mario Zanon
In multiclass softmax cross-
entropy, the full generalized Gauss-Newton (GGN)
curvature couples all output logits through the softmax covariance, making
curvature-vector products harder to scale as the number of classes grows. We show that the standard multiclass GGN can be decomposed exactly into a true-vs-rest term and a positive semidefinite within-competitor covariance term. Fast Gauss-Newton (FGN) retains the first term and drops the second, yielding a positive semidefinite under-approximation of the multiclass GGN that is exact for binary classification. The derivation uses an exact true-vs-rest scalar-margin representation of softmax cross-
entropy: the loss and gradient are unchanged, and the approximation enters only at the
curvature level. Exploiting the FGN
curvature structure, the damped update can be written as an equivalent whitened row-space system with one row per mini-batch example. We solve this system
matrix-free by conjugate gradient using Jacobian-vector and vector-Jacobian products of the scalar margin map. Targeted mechanism experiments show how the dropped covariance grows with competitor dispersion and how damping can reduce its effect on the resulting steps. In frozen-feature head optimization and end-to-end CIFAR-100 training, FGN achieves similar test accuracy to a full GGN baseline with faster updates.
Read the full abstract ▾
Machine Learning
Understanding diffusion models requires rethinking (again) generalization
Pierre Marion and Yu-Han Wu
This position paper argues that understanding generalization in
diffusion models requires fundamentally new theoretical frameworks that go beyond both classical statistical learning theory and the benign overfitting paradigm developed for supervised learning. In
diffusion models, unlike in supervised learning, memorization of training data and generalization to novel samples are incompatible: a model that has fully memorized its training set generates copies rather than novel data. Several theoretical explanations for why practical
diffusion models nevertheless generalize have been proposed, based on capacity limitations, implicit regularization from optimization, or architectural inductive biases, but their interactions remain unclear. We argue that the field should pivot from explaining why the
diffusion models do not memorize to investigating what the model actually learns during pre-memorization phase. To highlight our stance, we conduct empirical study of
diffusion models trained on CIFAR-10, and we distill the findings into concrete open questions that we believe are key to improve understanding of generalization in
diffusion models.
Read the full abstract ▾
Machine Learning
PRISM: Iterative Cross-Modal Posterior Refinement for Dynamic Text-Attributed Graphs
Trimble Chang, Yihang Liu, Mingjing Han et al. (+1 more)
Dynamic text-attributed graphs (DyTAGs) provide a powerful framework for modeling evolving systems in which node semantics and time-dependent interactions are tightly coupled. Recently,
multimodal learning has emerged as a promising yet underexplored direction for enhancing DyTAG representation learning. However, existing methods typically rely on rigid modality partitions and one-shot fusion strategies, which limit their ability to capture the intrinsic and evolving dependencies between node semantics and interaction behaviors. To address these limitations, we propose PRISM, an iterative cross-modal posterior refinement framework for DyTAG representation learning. PRISM organizes DyTAG information into semantic and behavioral modalities, providing a more intrinsic alternative to carrier-level modality partitions. Instead of fusing the two modalities in a single step, PRISM learns a refinement trajectory that progressively transforms semantic priors into behavior-conditioned posterior states through cross-modal interaction with behavioral evidence. Extensive experiments on DTGB
benchmark datasets show that PRISM achieves strong performance on temporal link prediction and destination node
retrieval tasks. Further ablation studies validate the effectiveness of semantic–behavioral modeling and iterative posterior refinement.
Read the full abstract ▾
All authors Trimble Chang, Yihang Liu, Mingjing Han, Han Zhang
Machine Learning
Normalized Architectures are Natively 4-Bit
Maxim Fishman, Brian Chmiel, Ron Banner et al. (+2 more)
Training large
language models at 4-bit precision is critical for efficiency. We show that nGPT, an architecture that constrains weights and hidden representations to the unit hypersphere, is inherently more robust to low-precision arithmetic. This removes the need for interventions-such as applying random Hadamard transforms and performing per-tensor scaling calculations-to preserve model quality, and it enables stable end-to-end NVFP4 training. We validate this approach on both a 1.2B dense model and hybrid (Mamba-
Transformer) MoE models of up to 3B/30B parameters. We trace this robustness to the dot product: while quantization noise remains largely uncorrelated in both standard and normalized architectures, the signal behaves differently. In nGPT, the hypersphere constraint enhances weak positive correlations among the element-wise products, leading to a constructive accumulation of the signal across the hidden dimension while the noise continues to average out. This yields a higher effective signal-to-noise ratio and a flatter loss landscape, with the effect strengthening as the hidden dimension grows, suggesting increasing advantages at scale. A reference implementation is available at https://github.com/anonymous452026/ngpt-nvfp4
Read the full abstract ▾
All authors Maxim Fishman, Brian Chmiel, Ron Banner, Daniel Soudry, Boris Ginsburg
Machine Learning
Causal Reinforcement Learning for Complex Card Games: A Magic The Gathering Benchmark
Cristiano da Costa Cunha, Ajmal Mian, Tim French et al. (+1 more)
Causal reinforcement learning (RL) lacks
benchmarks for complex systems that combine sequential decision making, hidden information, large masked action spaces, and explicit causal structure. We introduce MTG-Causal-RL, a Gymnasium
benchmark built on Magic: The Gathering with a 3,077-dimensional partial observation, a 478-action masked discrete action space, five competitive Standard archetypes, three reward schemes, and a hand-specified Structural Causal Model (SCM) over strategic variables. Every episode exposes causal variables, SCM-predicted intervention effects, and per-factor credit traces, making causal credit assignment, leave-one-out cross-archetype transfer, and policy auditability first-class metrics. We adapt a panel of reference baselines: random, heuristic, masked PPO, a causal-world-model PPO variant, and an architecture-matched scalar control. We propose Causal Graph-Factored Advantage PPO (CGFA-PPO) as a reference causal
agent that uses SCM parents of win probability as factor-aligned critic targets with an intervention-calibration loss. All comparisons use paired seeds, paired-bootstrap confidence intervals, and Holm-Bonferroni correction within pre-registered families. Masked PPO and CGFA-PPO reach competitive in-distribution win rates and exceed the random baseline; per-factor calibration trajectories and leave-one-out transfer gaps expose diagnostic structure that scalar win rate alone cannot. We release the
benchmark, reference-baseline results, and full evaluation protocol openly. By coupling a strategically rich, partially observed domain with an explicit causal interface and statistical protocol, MTG-Causal-RL gives causal-RL, world-model, and LLM-
agent research a shared testbed for questions current
benchmarks cannot pose together: causal credit assignment under masked action spaces, structural transfer across archetypes, and SCM-grounded policy auditability.
Read the full abstract ▾
All authors Cristiano da Costa Cunha, Ajmal Mian, Tim French, Wei Liu
Machine Learning
Geometry-Aware Simplicial Message Passing
Elena Xinyi Wang and Bastian Rieck
The Weisfeiler–Lehman (WL) test and its simplicial extension (SWL) characterize the combinatorial expressivity of message passing networks, but they are blind to geometry, i.e., meshes with identical connectivity but different
embeddings are indistinguishable. We introduce the
Geometric Simplicial Weisfeiler–Lehman (GSWL) test, which incorporates vertex coordinates into color refinement for
geometric simplicial complexes. In addition, we show that (i) the expressivity of geometry-aware simplicial message passing schemes is bounded above by GSWL, and (ii) that there exist parameters such that the discriminating power of GSWL is matched by these schemes on any fixed finite family of
geometric simplicial complexes. Combined with the Euler Characteristic Transform (ECT), a complete invariant for
geometric simplicial complexes, this yields a
geometric expressivity characterization together with an approximation framework. Experiments on synthetic and mesh
datasets serve to validate our theory, showing a clear hierarchy from combinatorial to geometry-aware models.
Read the full abstract ▾
Machine Learning · Cross-listed from stat.AP
Correcting heterogeneous diagnostic bias when developing clinical prediction models using causal hidden Markov models
Jose Benitez-Aurioles, Ricardo Silva, Brian McMillan et al. (+1 more)
In routine care, individuals identified a priori as high-risk are usually tested for conditions more frequently. Protected attributes, such as sex or ethnicity may also determine testing frequency. Such heterogeneous detection rates across a population induce label error. This causes systematic model error for specific groups and biases performance metrics during validation. This paper proposes a method to correct for such bias in prediction models due to differential diagnostic delay. We use a causal inference framework to define our target estimand: an individual's diagnosis probability in a counterfactual scenario where their diagnosis rate matches that of a reference group. We model the longitudinal process as a hidden Markov model, in which confirmatory test results are emissions from a latent progressive disease stage. We validate our approach in simulated data and apply it to a case study of chronic kidney disease prediction using electronic health records. In simulations, our method reduces prediction bias and improves calibration-in-the-large, correcting the Observed:Expected ratio in the underdiagnosed group from 1.34 (standard deviation: 0.09) in a model developed without any correction for underdiagnosis bias to 1.02 (0.09). Violations of assumptions in the simulation affected the estimation of model parameters, but the proposed approach nonetheless remained better calibrated than the standard model. In the clinical case study, we identify diabetes as the main driver of observability, with an odds ratio of 10.36 (95% confidence interval, 9.80 - 11.02) in 6-month urine albumin-creatinine ratio testing rate. Using our approach to predict the counterfactual diagnostic rate in patients without diabetes, we improved the Observed:Expected ratio of a developed clinical prediction model from 1.55 (1.51 - 1.59) to 1.01 (0.98 - 1.04).
Read the full abstract ▾
All authors Jose Benitez-Aurioles, Ricardo Silva, Brian McMillan, Matthew Sperrin
Machine Learning
Towards Self-Explainable Document Visual Question Answering with Chain-of-Explanation Predictions
Kjetil Indrehus, Adrian Duric, Changkyu Choi et al. (+1 more)
Document Visual Question Answering (DocVQA) requires
vision-
language models to reason not only about what information in a document is relevant to a question, but also where the answer is grounded on the page. Existing DocVQA models entangle question-relevant evidence and answer localization and operate largely as black boxes, offering limited means to verify how predictions depend on visual evidence. We propose CoExVQA, a self-explainable DocVQA framework with a grounded
reasoning process through a chain-of-explanation design. CoExVQA first identifies question-relevant evidence, then explicitly localizes the answer region, and finally decodes the answer exclusively from the grounded region. Prediction via CoExVQA's chain-of-explanation enables direct inspection and verification of the
reasoning process across modalities. Empirical results show that restricting decoding to grounded evidence achieves SotA explainable DocVQA performance on PFL-DocVQA, improving ANLS by 12% over the current explainable baselines while providing transparent and verifiable predictions.
Read the full abstract ▾
All authors Kjetil Indrehus, Adrian Duric, Changkyu Choi, Ali Ramezani-Kebrya
Machine Learning · Cross-listed from cs.DC
Relay Buffer Independent Communication over Pooled HBM for Efficient MoE Inference on Ascend
Tianlun Hu, Tiancheng Hu, Shengsheng Litang et al. (+8 more)
Mixture-of-Experts (MoE) inference requires large-scale
token exchange across devices, making dispatch and combine major bottlenecks in both prefill and decode. Beyond network transfer, routing-driven layout transformation, temporary relay, and output restoration can add substantial overhead. Existing MoE communication paths are often buffer-centric, using explicit inter-process relay and reordering buffers around collective transfer. This report presents a relay-buffer-free communication design for MoE inference acceleration on Ascend systems. The design reorganizes dispatch and combine around direct placement into destination expert windows and direct reading from remote expert windows. Built on globally pooled high-bandwidth
memory and symmetric-
memory allocation, it removes most intermediate relay and reordering buffers while retaining only lightweight control state, including counts, offsets, and synchronization metadata. We instantiate the design as two schedules for the main phases of MoE inference: a prefill schedule with richer planning state for throughput-oriented execution, and a compact decode schedule for latency-sensitive execution. Experiments on Ascend-based MoE workloads show reduced dispatch and combine latency in both settings. At the serving level, the implementation improves time to first
token (TTFT), preserves competitive time per output
token (TPOT), and enlarges the feasible scheduling space under practical latency constraints. These results indicate that, on platforms with globally addressable device
memory, reducing intermediate buffering and output restoration around expert execution is an effective direction for accelerating MoE inference.
Read the full abstract ▾
All authors Tianlun Hu, Tiancheng Hu, Shengsheng Litang, Sheng Wang, Xiaoming Bao, Yuxing Li, Wei Wang, Zhongzhe Hu, Lijun Li, Hongwei Sun, Jingbin Zhou
Machine Learning
Towards Generation-Efficient Uncertainty Estimation in Large Language Models
Mingcheng Zhu, Yu Liu and Tingting Zhu
Uncertainty estimation is important for deploying LLMs in high-stakes applications such as healthcare and finance, where
hallucinations can appear fluent and plausible while being factually incorrect, making it difficult for users to judge whether an output should be trusted. Existing methods require one or more full autoregressive generations to estimate uncertainty, which introduces substantial inference cost and often delays uncertainty assessment. In this paper, we investigate whether effective uncertainty estimation can be achieved with partial generation or even input-only information. Specifically, we first develop a unified framework that formulates uncertainty estimation as an early estimation problem over the autoregressive generation process of LLMs. This framework organises existing and proposed
estimators by the information they observe, ranging from multi-generation to input-only prediction, and clarifies the performance-cost trade-off underlying different uncertainty estimation methods. Building on this view, we study two largely underexplored low-cost settings: estimating uncertainty with part of the generation, and predicting uncertainty from the input
prompt. We propose Logit Magnitude, which uses top-M logit evidence to estimate uncertainty from an early-stopped generation prefix, and MetaUE, which distils generation-based uncertainty into a lightweight input-only
estimator trained with uncertainty scores. Extensive experiments on general and domain-specific
benchmarks show that Logit Magnitude achieves strong performance, and partial generations of LLMs are often sufficient for effective uncertainty estimation. MetaUE further provides a competitive input-only approximation in several settings. These findings suggest that effective uncertainty estimation requires less generation than commonly assumed, enabling unreliable responses to be identified earlier.
Read the full abstract ▾
Machine Learning
Learning beyond Site Bias for OOD Generalization in Brain Networks
Yingxu Wang, Kunyu Zhang, Yanwu Yang et al. (+5 more)
Graph-based learning from functional
magnetic resonance imaging (fMRI) has shown strong potential for
brain network analysis. However, existing methods often degrade under cross-site out-of-distribution (OOD) settings, as site-dependent confounder effects can obscure disease-related connectivity patterns and static functional connectivity (FC) does not explicitly capture informative within-scan variations. In this paper, we propose Cross-site OOD Robust
brain nEtwork (CORE), a unified framework for
brain network learning across unseen sites. First, CORE estimates site-specific confounder effects and aggregates the resulting deconfounders by cross-source reliability for unseen-site correction, while extracting a population scaffold of reproducible label-associated connections from site-wise residualized FC. It then summarizes temporal variations on scaffold edges into compact descriptors for line-graph modeling. Finally, prior-guided subject-adaptive gating modulates message passing, balancing population priors with individual variability. Extensive leave-one-site-out experiments on ABIDE, REST-meta-MDD, SRPBS, and ABCD demonstrate that CORE consistently outperforms competitive baselines, with up to a 10.3% relative improvement in accuracy. These gains also persist across different
brain parcellation schemes on ABIDE.
Read the full abstract ▾
All authors Yingxu Wang, Kunyu Zhang, Yanwu Yang, Mengzhu Wang, Thomas Wolfers, Yujie Wu, Siyang Gao, Nan Yin
Machine Learning
TFM-Retouche: A Lightweight Input-Space Adapter for Tabular Foundation Models
Duong Nguyen, Mohammed Jawhar and Nicolas Chesneau
Tabular foundation models (TFMs), such as TabPFN-2.6, TabICLv2, ConTextTab, Mitra, LimiX, and TabDPT, achieve strong zero-shot performance through in-context learning, but their inductive biases remain fixed at inference time. Adapting a pretrained TFM to a specific
dataset or task typically requires either full fine-tuning, which is computationally expensive, or parameter-efficient tuning methods (PEFT) such as LoRA, which must be tailored to the internal architecture of each TFM. Furthermore, the evidence on whether weight-space fine-tuning improves accuracy or calibration is mixed \citep{tanna_exploring_2026,rubachev_finetuning_2025}. We introduce TFM-Retouche, a lightweight input-space residual adapter that is architecture-agnostic by design with respect to the frozen TFM backbone. TFM-Retouche learns a small residual correction in the input space to align the input data with the inductive biases of the pretrained model. The adapter is trained end-to-end through the frozen TFM, with a post-training identity guard that falls back to the unmodified TFM whenever adaptation does not help on held-out validation. On TabArena-Lite (51
datasets spanning binary classification, multiclass classification, and regression), TabICLv2-Retouche – the framework instantiated on TabICLv2 – is the top-ranked method on the leaderboard with light per-task tuning and ensembling, lifting aggregate Elo by +56 over the frozen TabICLv2 base and sitting on the Pareto front of predictive quality versus both training and inference time.
Read the full abstract ▾
Machine Learning
Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference
Saksham Rathi, Preeti and Mythili Vutukuru
Auto-regressive
token generation in large
language models is
memory-bound because it requires "attending to" key and value tensors (KV cache) of all previous
tokens. Prior work aims to improve the efficiency of this decode process by batching multiple requests together, and maximizing batch size subject to GPU
memory constraints. The key observation of our work is that with prefix-sharing workloads, smaller, prefix-homogeneous batches - where all requests share a common prefix - can achieve higher decode throughput than larger, heterogeneous batches, due to better spatial and temporal locality during KV cache accesses. However, current prefix-aware schedulers maximize prefix reuse within a batch only to reduce KV cache
memory footprint, but do not stop batch formation at smaller homogeneous batches that could have performed better. Further, we show that shared prefix detection in existing schedulers relies on radix-tree traversals, incurring substantial CPU overhead that is often comparable to GPU execution time. This paper presents Feather, a prefix-aware scheduler that uses reinforcement learning (RL) to learn the optimal tradeoff between batch size and prefix homogeneity. We also introduce Chunked Hash Tree (CHT), a lightweight data structure that enables fast prefix detection and efficient request selection for the RL scheduler, avoiding expensive tree traversals. We integrate Feather into vLLM and SGLang, and our evaluation shows that Feather achieves 2-10$\times$ higher end-to-end throughput as compared to existing schedulers by reducing the total number of KV cache accesses, while doing no worse than the status quo when there is insufficient prefix sharing.
Read the full abstract ▾
Machine Learning
Optimal Transport for LLM Reward Modeling from Noisy Preference
Licheng Pan, Haochen Yang, Haoxuan Li et al. (+8 more)
Reward models are fundamental to Reinforcement Learning from Human Feedback (RLHF), yet real-world
datasets are inevitably corrupted by noisy preference. Conventional training objectives tend to overfit these errors, while existing denoising approaches often rely on homogeneous noise assumptions that fail to capture the complexity of linguistic preferences. To handle these challenges, we propose SelectiveRM, a framework grounded in optimal transport. We first devise a Joint Consistency Discrepancy to align the distribution of model predictions with preference data. Furthermore, to address the limitation of strict mass conservation which compels the model to fit outliers, we incorporate a Mass Relaxation mechanism via partial transport. This enables the autonomous exclusion of samples with noisy preference that contradict semantic consistency. Theoretically, we demonstrate that SelectiveRM optimizes a tighter upper bound on the unobserved clean risk. Extensive experiments validate that our approach significantly outperforms state-of-the-art baselines across diverse
benchmarks.
Read the full abstract ▾
All authors Licheng Pan, Haochen Yang, Haoxuan Li, Yunsheng Lu, Yongqi Tong, Yinuo Wang, Shijian Wang, Zhixuan Chu, Lei Shen, Yuan Lu, Hao Wang
Machine Learning
Does Synthetic Data Help? Empirical Evidence from Deep Learning Time Series Forecasters
Hugo Cazaux, Eyjólfur Ingi Ásgeirsson and Hlynur Stefánsson
Synthetic data has transformed
language model training, yet its role in time series
forecasting remains poorly understood. We present a large-scale empirical study: nine experiment groups, 4,218 runs systematically evaluating synthetic time series augmentation across five architectures, four synthetic signals and seven
datasets. The effect is sharply architecture-conditional: channel-mixing models (TimesNet, iTransformer) benefit in the majority of trials, while channel-independent models (DLinear, PatchTST) are consistently degraded. In selected low-resource settings the gains are striking: TimesNet trained on only 10% of Weather data with synthetic augmentation surpasses the full-data baseline (4 of 16 sparsity-
dataset combinations). Averaged across all architectures, augmentation hurts in 67% of trials. We further find that only the Seasonal-Trend generator reliably helps across the tested
benchmarks, and that hard curriculum switching is actively harmful (+24% MSE degradation). These results provide concrete, actionable guidelines on how to use synthetic data: use synthetic augmentation with channel-mixing architectures, use gradual annealing schedules, and treat low-resource augmentation as architecture- and
dataset-dependent. Code is available at \href{https://github.com/hugoiscracked/synthetic-ts/tree/main}
Read the full abstract ▾
Machine Learning
Multi-agent decision making: A Blackwell's informativeness approach
Zheng Zhang, Cuong C. Nguyen, Kevin Wells et al. (+1 more)
The rapid development of large
language models (LLMs) has motivated research on decision-making in multi-
agent systems, where multiple
agents collaborate to achieve shared objectives. Existing aggregation approaches, such as voting and debate, are largely ad-hoc and lack formal guarantees regarding the informativeness of the resulting decisions. In this paper, we provide a principled approach to analyse decisions made in the multi-LLM setting using Blackwell's informativeness framework. Within the Blackwell information-structure abstraction, we show that voting and debate induce information structures that are no more informative than the pooled private information of all
agents. This result identifies
Bayesian pooled posterior maximisation as an information-theoretic upper-bound decision rule under the Blackwell ordering. Motivated by this theoretical analysis, we introduce a practical method for LLM-based question-answering (QA) tasks that estimates each
agent's posterior and approximates the pooled posterior using a product-of-posteriors
estimator. Extensive experiments on six QA
benchmarks demonstrate that our approach outperforms state-of-the-art multi-LLM debate and voting methods.
Read the full abstract ▾
All authors Zheng Zhang, Cuong C. Nguyen, Kevin Wells, Gustavo Carneiro
Machine Learning
Matrix-Decoupled Concentration for Autoregressive Sequences: Dimension-Free Guarantees for Sparse Long-Context Rewards
Pei-Sen Li
Sequence-level evaluations in autoregressive Large
Language Models (LLMs) rely on highly dependent
token generation. Establishing tight concentration bounds for these processes remains a challenge due to two fundamental bottlenecks in existing frameworks: (i) classical inequalities typically separate dependency structures from target sensitivities, leading to a scalar collapse that inflates the variance proxy to a suboptimal $\mathcal{O}(N)$ for sparse terminal rewards; (ii) conversely, while certain spatial methods achieve tighter bounds, they lack the strictly causal filtration required by sequential generation, rendering them inapplicable to the autoregressive setting. To resolve both bottlenecks, we establish a sharp McDiarmid-type inequality for dependent sequences, governed strictly by the exact
matrix-vector multiplication of the causal dependency resolvent and the target sensitivity vector. This
Matrix-Decoupled Concentration (MDC) framework natively recovers optimal constants for Markov chains and exploits directed $d$-separation to yield order-optimal bounds for causal trees. Crucially, by exactly preserving the coordinate-wise sparsity of rewards within a strictly causal framework, MDC mathematically prevents scalar collapse, guaranteeing a dimension-free $\mathcal{O}(1)$ variance proxy and providing a rigorous mathematical justification for the stability of long-context
reasoning.
Read the full abstract ▾
Machine Learning
Quantizing With Randomized Hadamard Transforms: Efficient Heuristic Now Proven
Ran Ben-Basat, William Kuszmaul, Michael Mitzenmacher et al. (+2 more)
Uniform random rotations (URRs) are a common preprocessing step in modern quantization approaches used for gradient compression, inference acceleration, KV-cache compression, model weight quantization, and approximate nearest-neighbor search in vector databases. In practice, URRs are often replaced by randomized Hadamard transforms (RHTs), which preserve orthogonality while admitting fast implementations. The remaining issue is the performance for worst-case inputs. With a URR, each coordinate is distributed as a shifted beta distribution, which converges to a Gaussian distribution in high dimensions. Generally, one RHT is not suitable in the worst case, as individual coordinates can be far from these distributions. We show that after composing two RHTs on any $d$-sized input vector, the marginal distribution of every fixed coordinate of the normalized rotated vector is within $O(d^{-1/2})$ of a standard Gaussian both in Kolmogorov distance and in $1$-Wasserstein distance. We then plug these bounds into the analyses of modern compression schemes, namely DRIVE and QUIC-FL, and show that two RHTs achieve performance that asymptotically matches URRs. However, while two RHTs suffice for scalar quantization, they may be insufficient for Vector Quantization (VQ), which depends on the joint distribution across fixed-size blocks of coordinates (as opposed to only marginal distribution convergence for single coordinates). We prove that a composition of three RHTs leads to decaying coordinate covariance. This ensures that any fixed, bounded, multi-dimensional VQ codebook optimized for URRs has the same expected error when using three RHTs, up to an additive term that vanishes with the dimension. Finally, because practical inputs are rarely
adversarial, we propose a linear-time $O(d)$ check on the input's moments to dynamically adapt the number of RHTs used at runtime to improve performance.
Read the full abstract ▾
All authors Ran Ben-Basat, William Kuszmaul, Michael Mitzenmacher, Amit Portnoy, Shay Vargaftik