Machine Learning
Do Fair Models Reason Fairly? Counterfactual Explanation Consistency for Procedural Fairness in Credit Decisions
Gideon Popoola and John Sheppard
Machine learning
algorithms in socially sensitive domains (e.g., credit decisions) often focus on equalizing predictive outcomes. However, satisfying these metrics does not guarantee that models use the same
reasoning for different groups. We show that existing outcome-fair models can still apply fundamentally different
reasoning to individuals, a “hidden procedural bias” missed by standard
fairness metrics and
algorithms. We propose Counterfactual Explanation Consistency (CEC), a framework that detects and mitigates this bias by aligning feature attributions between individuals and their counterfactual counterparts. Key contributions include a nearest-neighbor counterfactual generation method, a modified baseline for integrated gradient comparisons, an individual-level procedural
fairness metric, and a corresponding training loss. We introduce a taxonomy identifying “Regime B” (same outcome, different
reasoning) as a critical blind spot. Experiments on synthetic data, German Credit, Adult Income, and HMDA mortgage data demonstrate that outcome-fair baselines exhibit substantial hidden bias, while CEC substantially reduces it with modest utility cost.
Read the full abstract ▾
Machine Learning
UFO: A Domain-Unification-Free Operator Framework for Generalized Operator Learning
Hanli Qiao, George Em Karniadakis and Muhammad Muniruzzaman
Neural operators have become an effective framework for learning mappings between function spaces, yet most existing architectures realize operators within a single representational domain, such as physical, spectral, or latent space. In this work, we introduce UFO (Domain-Unification-Free Operator), a cross-domain neural operator framework that realizes operators through adaptive, jointly conditioned interactions among representations defined on distinct domains. UFO enables discretization decoupling: the input function can be observed at resolutions or locations different from those used during training, while the solution can be queried at arbitrary output resolutions. Across four complementary
benchmarks covering discontinuous inputs, irregular sampling with spectral mismatch, nonlinear dynamics, and
stochastic high-frequency fields, UFO delivers accurate, robust, and physically coherent predictions under distribution shifts. These results establish cross-domain, phase-modulated realization as a powerful framework for discretization-decoupled neural operator learning.
Read the full abstract ▾
Machine Learning
Modeling Heterophily in Multiplex Graphs: An Adaptive Approach for Node Classification
Kamel Abdous, Nairouz Mrabah and Mohamed Bouguessa
Existing multiplex graph models often assume homophily, where connected nodes tend to belong to the same class or share similar attributes. Consequently, these models may struggle with graphs exhibiting heterophily, where connected nodes typically belong to different classes and have dissimilar attributes. While recent methods have been developed to learn reliable node representations from unidimensional graphs with heterophily, they do not fully address the complexities of multiplex graphs. In a multiplex graph, nodes are linked through multiple types of edges (referred to as dimensions), which can simultaneously exhibit homophilic and heterophilic interactions. To address this gap, we propose \methodname, a novel method for node classification in multiplex graphs that adapts to both homophilic and heterophilic dimensions. \methodname introduces dimension-specific compatibility
matrices to model varying degrees of homophily and heterophily across dimensions. A key innovation is its use of a product of trainable low-pass and high-pass filters, approximated via Chebyshev
polynomials, to capture both smooth and abrupt changes in the graph signal. By composing these filters and optimizing label predictions using a proximal-gradient method, \methodname dynamically adjusts to the heterophilic characteristics of each dimension. Extensive experiments on synthetic and real-world
datasets provide evidence that \methodname captures the complex interplay of homophilic and heterophilic interactions in multiplex graphs, and tends to yield improved node classification performance compared to state-of-the-art methods.
Read the full abstract ▾
Machine Learning · Cross-listed from stat.ML
A Unified Framework for Critical Scaling of Inverse Temperature in Self-Attention
Tomohiro Hayase and Ryo Karakida
Length-dependent logit rescaling is widely used to stabilize long-context self-
attention, but existing analyses and methods suggest conflicting inverse-temperature laws for the context length $n$, ranging from $(\log n)^{1/2}$ to $\log n$ and $(\log n)^2$. We provide a general theory showing that the desirable scale is determined by the gap-counting function $N_n$ of each
attention row. Counting how many competitors lie within each gap from the maximum, we define an upper-tail accumulation scale and prove that it gives the critical inverse-temperature scale for softmax concentration: below this scale, the top competitors remain unseparated, whereas above it, the
attention entropy collapses. This framework unifies prior scaling laws as different $N_n$ and yields a direct diagnostic for
attention-score families, from idealized theoretical models to more practical
transformers.
Read the full abstract ▾
Machine Learning
IGT-OMD: Implicit Gradient Transport for Decision-Focused Learning under Delayed Feedback
Benjamin Amoh, Geoffrey G. Parker and Wesley Marrero
Decision-focused learning trains predictive models end-to-end against downstream decision loss, but online settings suffer delayed feedback: outcomes may not arrive for many environment interactions. We identify staleness amplification, a failure mode unique to bilevel optimization under delay, in which gradient staleness couples with inner-solver sensitivity to inflate regret beyond single-level delay theory. We prove that any black-box delayed optimizer incurs an irreducible regret cost from inner-solver approximation error, and that gradient staleness contributes a quadratically growing transport error without bilevel-aware correction. Our
algorithm, IGT-OMD, applies Implicit Gradient Transport to hypergradients within Online Mirror Descent, re-evaluating stale gradients at the current parameters using stored inner solutions. This method reduces transport error from a quadratic to a linear dependence on delay and achieves the first sublinear regret bound for delayed bilevel optimization with queue-length-adaptive step sizes. Controlled experiments provide a mechanistic fingerprint: transport benefit is exactly $0.0\%$ ($p=1.00$) at unit delay and grows monotonically to $9.5\%$ at fifty rounds ($p<0.001$), isolating the correction's effect. On Linear Quadratic Regulator, Warcraft shortest-path, and Sinkhorn optimal transport, IGT-OMD reduces decision loss by $17$–$55\%$ relative to single-level baselines, with phase transitions matching the theory.
Read the full abstract ▾
Machine Learning
A Unified Perspective for Learning Graph Representations Across Multi-Level Abstractions
Mohamed Mahmoud Amar, Nairouz Mrabah, Mohamed Bouguessa et al. (+1 more)
Graph Self-Supervised Learning (GSSL) has emerged as a powerful paradigm for generating high-quality representations for graph-structured data. While multi-scale graph contrastive learning has received increasing
attention, many existing methods still predominantly focus on a single graph abstraction level. To address this limitation, we propose a unified contrastive framework that can target node-level, proximity-level, cluster-level, and graph-level information and integrate them through a linear combination of similarity scores on positive pairs and dissimilarity scores (i.e., similarity scores on negative pairs). Furthermore, current approaches typically assign uniform penalty strengths to all examples, which reduces optimization flexibility and leads to ambiguous convergence status. To overcome this, we introduce a novel parameter-free fine-grained self-weighting mechanism that adaptively assigns weights to individual similarity and dissimilarity scores. The proposed mechanism emphasizes the scores that deviate significantly from their target values. Our approach not only enhances optimization flexibility but also eliminates the computational overhead of hyperparameter tuning in conventional multi-task GSSL methods. Comprehensive experiments on real-world
datasets show that our methods consistently outperform state-of-the-art approaches across downstream tasks, including classification, clustering, and link prediction, in both single-level and multi-level scenarios.
Read the full abstract ▾
All authors Mohamed Mahmoud Amar, Nairouz Mrabah, Mohamed Bouguessa, Abdoulaye Baniré Diallo
Machine Learning
Parallel-in-Time Training of Recurrent Neural Networks for Dynamical Systems Reconstruction
Florian Hess, Florian Götz and Daniel Durstewitz
Reconstructing nonlinear dynamical systems (DS) from data (DSR) is a fundamental challenge in science and engineering, but it inherently relies on sequential models. Recent breakthroughs for sequential models have produced
algorithms that parallelize computation along sequence length $T$, achieving logarithmic time complexity, $\mathcal{O}(\log T)$. Since sequence lengths have been practically limited due to the linear runtime complexity $\mathcal{O}(T)$ of classical backpropagation through time, this opens new avenues for DSR. This paper studies two prominent classes of parallel-in-time
algorithms for this task, both of which leverage parallel associative scans as their core computational primitive. The first class comprises models with linear yet non-autonomous dynamics and a nonlinear readout, such as modern State Space Models (SSMs), while the second consists of general nonlinear models which can be parallelized using the DEER framework. We find that the linear training-time recurrence of the first class of models imposes limitations that often hinder learning of accurate nonlinear dynamics. To address this, we augment DEER with Generalized Teacher Forcing (GTF), a novel variant within the more general nonlinear framework that ensures stable and effective learning of nonlinear dynamics across arbitrary sequence lengths. Using GTF-DEER, we investigate the benefits of training on extremely long sequences ($T$ approaching $10^5$) for DSR. Our results show that access to such long trajectories significantly improves DSR if the data features long time scales. This work establishes GTF-DEER as a robust tool for data-driven discovery and underscores the largely untapped potential of long-sequence learning in modeling complex DS.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.AI
Revealing Interpretable Failure Modes of VLMs
Isha Chaudhary, Vedaant V Jain, Kavya Sachdeva et al. (+2 more)
Vision-
Language Models (VLMs) are increasingly used in safety-critical applications because of their broad
reasoning capabilities and ability to generalize with minimal task-specific engineering. Despite these advantages, they can exhibit catastrophic failures in specific real-world situations, constituting failure modes. We introduce REVELIO, a framework for systematically uncovering interpretable failure modes in VLMs. We define a failure mode as a composition of interpretable, domain-relevant concepts-such as pedestrian proximity or adverse weather conditions-under which a target VLM consistently behaves incorrectly. Identifying such failures requires searching over an exponentially large discrete combinatorial space. To address this challenge, REVELIO combines two search procedures: a diversity-aware beam search that efficiently maps the failure landscape, and a Gaussian-process Thompson Sampling strategy that enables broader exploration of complex failure modes. We apply REVELIO to autonomous driving and indoor robotics domains, uncovering previously unreported vulnerabilities in state-of-the-art VLMs. In driving environments, the models often demonstrate weak spatial grounding and fail to account for major obstructions, leading to recommendations that would result in simulated crashes. In indoor robotics tasks, VLMs either miss safety hazards or behave excessively conservatively, producing false alarms and reducing operational efficiency. By identifying structured and interpretable failure modes, REVELIO offers actionable insights that can support targeted VLM safety improvements.
Read the full abstract ▾
All authors Isha Chaudhary, Vedaant V Jain, Kavya Sachdeva, Sayan Ranu, Gagandeep Singh
Machine Learning · Cross-listed from stat.ML
Online Conformal Prediction: Enforcing monotonicity via Online Optimization
Eduardo Ochoa Rivera and Ambuj Tewari
Conformal prediction provides a principled framework for uncertainty quantification with finite-sample coverage guarantees. While recent work has extended conformal prediction to online and sequential settings, existing methods typically focus on a single coverage level and do not ensure consistency across multiple confidence levels. In many real-world applications, such as weather
forecasting, macroeconomic prediction, and risk management, different users operate under heterogeneous risk tolerances and require calibrated uncertainty estimates across a range of coverage levels. In such settings, it is desirable to produce prediction sets corresponding to different coverage levels that are nested and valid simultaneously. In this paper, we propose two novel online conformal prediction methods that output nested prediction sets across a range of coverage levels, enabling simultaneous uncertainty quantification across the entire risk spectrum. Beyond interpretability, jointly estimating multiple coverage levels is known to improve statistical efficiency in classical
quantile regression by enforcing non-crossing constraints and sharing information across
quantiles. Our approaches leverage an online optimization perspective with small regret that translates to
quantile estimation error control while enforcing nestedness of prediction sets. Empirical results on synthetic and real-world
datasets, including applications in
forecasting tasks with heterogeneous risk requirements, demonstrate that our method achieves stable coverage across all levels, strictly nested prediction sets, and improved efficiency compared to existing online conformal baselines.
Read the full abstract ▾
Machine Learning
ODRPO: Ordinal Decompositions of Discrete Rewards for Robust Policy Optimization
Nirmal Patel, Fei Wang and Inderjit S. Dhillon
The
alignment of Large
Language Models (LLMs) utilizes Reinforcement Learning from AI Feedback (RLAIF) for non-verifiable domains such as long-form question answering and open-ended instruction following. These domains often rely on LLM based auto-raters to provide granular, multi-tier discrete rewards (e.g., 1-10 rubrics) that are inherently
stochastic due to
prompt sensitivity and sampling randomness. We empirically verify the stochasticity of auto-raters that can propagate and corrupt standard advantage
estimators like GRPO and MaxRL, as a noisy reward samples can skew normalization statistics and degrade the global learning signal. Empirically, sampling more rewards and taking majority voting may reduce the noise and improve performance, but this approach is computationally expensive. To address this bottleneck, we introduce $\textbf{O}$rdinal $\textbf{D}$ecomposition for $\textbf{R}$obust $\textbf{P}$olicy $\textbf{O}$ptimization ($\textbf{ODRPO}$), a framework that structurally isolates evaluation noise by decomposing discrete rewards into a sequence of ordinal binary indicators. By independently computing and accumulating advantages across these progressively challenging success thresholds, ODRPO prevents outlier evaluations from corrupting the global update while establishing an implicit, variance-aware learning curriculum. Empirically, ODRPO achieves robust performance on Qwen2.5-7B and Qwen3-4B models, outperforming baselines with relative improvements of upto 14.8% on FACTS-grounding-v2 and 7.5% on Alpaca-Evals. Critically, these gains are achieved with negligible training-time overhead, as ODRPO requires no additional compute per step compared to standard
estimators. Supported by theoretical analysis confirming its optimization stability, ODRPO provides a scalable and robust framework for aligning models within the noisy, discrete evaluation landscape of modern RLAIF.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.GT
Profit Maximization in Bilateral Trade against a Smooth Adversary
Simone Di Gregorio, Paul Dütting, Federico Fusco et al. (+1 more)
Bilateral trade models the task of intermediating between two strategic
agents, a seller and a buyer, who wish to trade a good. We study this problem from the perspective of a profit-maximizing broker within an online learning framework, where the
agents' valuations are generated by a smooth adversary. We devise a learning
algorithm that guarantees a $\tilde{O}(\sqrt{T})$ regret bound, which is tight in the time horizon $T$ up to poly-logarithmic factors. This matches the minimax rate for the
stochastic i.i.d. case, and is also well separated from the
adversarial setting, where sublinear-regret is unattainable. By extending the strong regret guarantees from the i.i.d. case to the smooth adversary, we significantly broaden the scope of settings where such fast rate is achievable, while closing an important gap in the regret landscape of this fundamental economic problem. To overcome the challenges posed by this adversary, we leverage a continuity property of smooth instances and combines this with a hierarchical net-construction of the broker's action space, which is analyzed via algorithmic chaining. We showcase the applicability of these techniques by deriving a similarly tight $\tilde{O}(\sqrt{T})$ regret bound for a related mechanism design model: the joint ads problem.
Read the full abstract ▾
All authors Simone Di Gregorio, Paul Dütting, Federico Fusco, Chris Schwiegelshohn
Machine Learning
scShapeBench: Discovering geometry from high dimensional scRNAseq data
Andrew J Steindl, João Felipe Rocha, Brian Tshilengi Di Bassinga et al. (+13 more)
High-dimensional point cloud data arise across many scientific domains, especially single-
cell biology. The shapes or
topologies of these
datasets determine the types of information that can be extracted. For example, clustered data supports
cell-type identification, trajectory structures support transition analysis, and archetypal structures capture continua of cellular behaviors. Existing analysis pipelines often assume a specific shape. The standard Seurat pipeline combines UMAP visualization with Louvain clustering and therefore assumes clustered data, while tools such as Monocle and SPADE assume tree-like structures, and flow-based models such as MIOFlow and Conditional Flow Matching target trajectories. Choosing which pipeline to apply is therefore often left to bioinformaticians who visually inspect
datasets before selecting an analysis strategy. With the rise of agentic AI scientists, automating shape detection is increasingly important for selecting downstream analysis pipelines. To address this problem, we introduce scShapeBench, a
benchmark dataset for shape detection containing both synthetic and expert-annotated single-
cell datasets. Synthetic
datasets are sampled from ground-truth skeleton graphs with controlled variance. Real single-
cell datasets are curated from diverse sources and annotated by experts into four categories: clusters, single trajectory, multi-branching, and archetypal. We additionally introduce scReebTower, a baseline method that uses
diffusion geometry to extract Reeb graphs and connect visualization with pipeline selection. We provide
topology-aware evaluation metrics and compare scReebTower against PAGA and Mapper on synthetic and real data. Our results indicate that scReebTower outperforms existing baselines. Overall, our contributions span
benchmarks, evaluation metrics, and a baseline for automated shape detection in single-
cell data.
Read the full abstract ▾
All authors Andrew J Steindl, João Felipe Rocha, Brian Tshilengi Di Bassinga, Zachary Warren, Matthew Scicluna, César Miguel Valdez Córdova, Shabarni Gupta, Leire Torices, Daniel Neumann, Timothy J. Mann, Ihuan Gunawan, Dhananjay Bhaskar, John G Lock, Christine L Chaffer, Guy Wolf, Smita Krishnaswamy
Machine Learning
Plan Before You Trade: Inference-Time Optimization for RL Trading Agents
Eun Go, Rohan Deb and Arindam Banerjee
Reinforcement learning
agents for portfolio management are typically trained and deployed as static policies, with no mechanism for policy adjustments using price forecasts at inference time. We propose FinPILOT (Financial Plugin Inference-time Learning for Optimal Trading), a plugin inference-time policy optimization framework inspired by Model Predictive Control (MPC). Under the price-taker assumption, portfolio actions do not affect future
market prices, allowing a multi-step price forecaster to provide imagined
market trajectories without iterative action-conditioned rollouts. At each decision step, FinPILOT uses the forecaster's predicted price trajectory to construct an allocation-based imagined return objective and optimizes the pretrained policy at inference time before executing one step of the trade. Our framework is compatible with any pretrained policy and adapts the policy using the forecaster's predictions without modifying the original RL training pipeline. Evaluated with five pretrained policies (PPO, SAC, A2C, TD3, DDPG) across DJ30 equities and foreign exchange (FX) over three rolling held-out test years (30 settings in total), FinPILOT improves Total Return in 28 of 30 settings and Sharpe ratio in 27 of 30, with gains spanning both
stochastic and deterministic policies. Direct single- and multi-period portfolio optimizers supplied with the same forecasts do not reproduce these aggregate gains under our evaluation protocol. Further, using controlled forecasts at calibrated quality levels, we show that both the magnitude and consistency of FinPILOT's performance increase as forecast error is reduced.
Read the full abstract ▾
Machine Learning
Multi-Rollout On-Policy Distillation via Peer Successes and Failures
Weichen Yu, Xiaomin Li, Yizhou Zhao et al. (+8 more)
Large
language models are often post-trained with sparse verifier rewards, which indicate whether a sampled trajectory succeeds but provide limited guidance about where
reasoning succeeds or fails. On-policy
distillation (OPD) offers denser
token-level supervision by training on student-generated trajectories, yet existing methods typically distill each rollout independently and ignore the other attempts sampled for the same
prompt. We introduce Multi-Rollout On-Policy
Distillation (MOPD), a peer-conditioned
distillation framework that uses the student's local rollout group to construct more informative teacher signals. MOPD conditions the teacher on both successful and failed peer rollouts: successes provide positive evidence for valid
reasoning patterns, while failures provide structured negative evidence about plausible mistakes to avoid. We study two peer-context constructions: positive peer imitation and contrastive success-failure conditioning. Experiments on competitive programming, mathematical
reasoning, scientific question answering, and tool-use
benchmarks show that MOPD consistently improves over standard on-policy baselines. Further teacher-signal analysis shows that mixed success-failure contexts better align teacher scores with verifier rewards, indicating that the gains arise from more faithful, instance-adaptive supervision. These results indicate that effective on-policy
distillation should exploit the student's multi-rollout trial-and-error behavior rather than treating rollouts as isolated samples.
Read the full abstract ▾
All authors Weichen Yu, Xiaomin Li, Yizhou Zhao, Xiaoze Liu, Ruowang Zhang, Haixin Wang, Yinyi Luo, Chen Henry Wu, Gaurav Mittal, Matt Fredrikson, Yu Hu
Machine Learning
Runtime Monitoring of Perception-Based Autonomous Systems via Embedding Temporal Logic
Parv Kapoor, Abigail Hammer, Ashish Kapoor et al. (+2 more)
Runtime monitoring of autonomous systems traditionally relies on mapping continuous sensor observations to discrete logical propositions defined over low-dimensional state variables. This abstraction breaks down in perception-driven settings, where such mappings require additional learned modules that are often computationally expensive, brittle, and semantically misaligned. In this work, we propose
Embedding Temporal Logic (ETL), a temporal logic that performs monitoring directly in learned
embedding spaces. ETL defines predicates through distances between observed
embeddings and target
embeddings derived from reference observations. This formulation allows specifications to capture high-level perceptual concepts, such as similarity to visual goals or avoidance of semantic regions, that are difficult or impossible to express using traditional predicates. By composing these predicates with temporal operators, ETL naturally expresses temporally extended and sequential perceptual behaviors. We introduce ETL monitors for evaluating specifications over bounded
embedding traces, along with a conformal calibration procedure that provides reliable and safety-oriented predicate evaluation. We evaluate our approach across multiple manipulation environments to show that ETL achieves strong empirical agreement with ground-truth semantics, including accurate monitoring of temporally composed behaviors.
Read the full abstract ▾
All authors Parv Kapoor, Abigail Hammer, Ashish Kapoor, Karen Leung, Eunsuk Kang
Machine Learning
Optimization Risk Bounds for Kolmogorov-Arnold Networks Trained by DP-SGD with Correlated Noise
Puyu Wang, Jan Schuchardt, Nikita Kalinin et al. (+5 more)
The theoretical understanding of differentially private
stochastic gradient descent (DP-SGD) with temporally correlated noise remains limited, particularly for non-convex neural network training. As a first step, we study two-layer Kolmogorov-Arnold Networks (KANs), a recently introduced architecture with learnable spline-based edge functions. We establish the first optimization risk bounds for clipped mini-batch DP-SGD with correlated noise in this setting, with explicit dependence on temporal correlation, clipping, mini-batch sampling, and network width. Existing arguments fail for three reasons: temporal dependence breaks the conditional-centering step; projection obstructs the cross-iteration cancellation of correlated perturbations; and active clipping breaks the empirical-gradient structure. We address these issues through shifted and auxiliary dynamics, a weighted empirical loss, and a high-probability localization argument. Our bound shows that temporal correlation reduces the leading noise terms, the clipping threshold enters essentially through an effective step size, and private training admits an explicit network width range. Experiments on synthetic data and MNIST support the predicted optimization effect of temporal correlation. As an application, we derive population risk guarantees via algorithmic stability. Our framework recovers non-private mini-batch SGD, independent-noise DP-SGD, and their full-batch counterparts as special cases.
Read the full abstract ▾
All authors Puyu Wang, Jan Schuchardt, Nikita Kalinin, Marcio Monteiro, Junyu Zhou, Sophie Fellenz, Christoph Lampert, Marius Kloft
Machine Learning
Learning to Decide with AI Assistance under Human-Alignment
Nina Corvelo Benz, Eleni Straitouri and Manuel Gomez-Rodriguez
It is widely agreed that when AI models assist decision-makers in high-stakes domains by predicting an outcome of interest, they should communicate the confidence of their predictions. However, empirical evidence suggests that decision-makers often struggle to determine when to trust a prediction based solely on this communicated confidence. In this context, recent theoretical and empirical work suggests a positive correlation between the utility of AI-assisted decision-making and the degree of
alignment between the AI confidence and the decision-makers' confidence in their own predictions. Crucially, these findings do not yet elucidate the extent to which this
alignment influences the complexity of learning to make optimal decisions through repeated interactions. In this paper, we address this question in the canonical case of binary predictions and binary decisions. We first show that this problem is equivalent to a two-armed online contextual learning problem with full feedback, and establish a lower bound of $\Omega (\sqrt{|H| \cdot |B| \cdot T} )$ on the expected regret any learner can attain, where $H$ and $B$ denote the sets of human and AI confidence values. We then demonstrate that, under perfect
alignment between AI and human confidence, a learner can attain an expected regret of $O(\sqrt{|H| \cdot T\log T})$ and, when $\sqrt{|H|} = O(\log T)$ and $B$ is countable, a non-trivial generalization of the Dvoretzky-Kiefer-Wolfowitz inequality improves the regret bound to $O(\sqrt{T\log T})$. Taken together, these results reveal that
alignment can reduce the complexity of learning to make decisions with AI assistance. Experiments on real data from two different human-subject studies where participants solve simple decision-making tasks assisted by AI models show that our theoretical results are robust to violations of perfect
alignment.
Read the full abstract ▾
Machine Learning
OceanCBM: A Concept Bottleneck Model for Mechanistic Interpretability in Ocean Forecasting
Sanah Suri, Kieran Ringel and Maike Sonnewald
Extreme ocean phenomena are challenging not only to predict but to diagnose, as accurate forecasts alone do not reveal the underlying physical drivers. While recent machine learning approaches achieve strong predictive skill, they remain largely opaque and provide limited guarantees of fidelity to ground-truth physics. We introduce OceanCBM, the first concept bottleneck model (CBM) for spatiotemporal prediction and mechanistic interrogation of ocean dynamics. OceanCBM uses mixed supervision to predict mixed layer heat content, a key precursor of marine heatwaves, while routing information through an intermediate layer of prescribed concepts derived from geophysical fluid dynamics and a 'free' concept. This design imposes soft physical structure without over-constraining the model, and the free concept both regularizes concept predictions and captures residual physical processes. Across ensemble initializations, we show that mixed supervision yields consistent mechanistic representations, whereas prediction-only and prescription-only baselines learn highly variable latent structures despite similar predictive performance. OceanCBM achieves interpretable, physically grounded representations without sacrificing skill, explicitly characterizing the interpretability-performance trade-off.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
DocAtlas: Multilingual Document Understanding Across 80+ Languages
Ahmed Heakl, Youssef Mohamed, Abdullah Sohail et al. (+6 more)
Multilingual document understanding remains limited for low-resource
languages due to scarce training data and model-based annotation pipelines that perpetuate existing biases. We introduce DocAtlas, a framework that constructs high-fidelity OCR
datasets and
benchmarks covering 82
languages and 9 evaluation tasks. Our dual pipelines, differential rendering of native DOCX documents and synthetic LaTeX-based generation for right-to-left scripts produce precise structural annotations in a unified DocTag format encoding layout, text, and component types, without learned models for core annotation. Evaluating 16 state-of-the-art models reveals persistent gaps in low-resource scripts. We show that Direct Preference Optimization (DPO) using rendering-derived ground truth as positive signal achieves stable
multilingual adaptation, improving both in-domain (+1.9%) and out-of-domain (+1.8%) accuracy without measurable base-
language degradation, where supervised fine-tuning degrades out-of-domain performance by up to 21%. Our best variant, DocAtlas-DeepSeek, improves +1.7% over the strongest baseline. Code is available at https://github.com/ahmedheakl/DocAtlas .
Read the full abstract ▾
All authors Ahmed Heakl, Youssef Mohamed, Abdullah Sohail, Rania Elbadry, Ahmed Nassar, Peter W. J. Staar, Fahad Shahbaz Khan, Imran Razzak, Salman Khan
Machine Learning · Cross-listed from cond-mat.dis-nn
The critical slowing down in training diffusion models
Luca Maria Del Bono, Giulio Biroli, Patrick Charbonneau et al. (+1 more)
Computational sampling has been central to the sciences since the mid-20th century. While machine-learning-based approaches have recently enabled major advances, their behavior remains poorly understood, with limited theoretical control over when and why they succeed. Here we provide such insight for
diffusion models—a class of generative schemes highly effective in practice—by analyzing their application to the $O(n)$ model of statistical field theory in the Gaussian limit $n \to \infty$. In this analytically tractable setting, we show that training a score model with a one-layer network architecture matching the exact solution exhibits a form of critical slowing down in parameter learning. This slowing down also impacts the generation process, indicating that the well-known difficulties of sampling near criticality persist even for learned generative models. To overcome this bottleneck, we consider the power of architectural depth. We find that using a two-layer architecture drastically reduces the critical slowing down, with the training time scaling logarithmically rather than quadratically with system size. Using a
Fourier implementation of the architecture, we further show that this acceleration in training time can be achieved without drastically increasing operational complexity. Taken together, these results demonstrate that
diffusion models can overcome the critical slowing down through appropriate architectural design, and establish a controlled framework for understanding and improving learned sampling methods in statistical physics and beyond.
Read the full abstract ▾
All authors Luca Maria Del Bono, Giulio Biroli, Patrick Charbonneau, Marylou Gabrié
Machine Learning
Towards Robust Federated Multimodal Graph Learning under Modality Heterogeneity
Sirui Zhang, Haonan Wang, Xunkai Li et al. (+5 more)
Recently,
multimodal graph learning (MGL) has garnered significant
attention for integrating diverse modality information and structured context to support various network applications. However, real-world graphs are often isolated due to data-sharing limitations across multiple parties, and their modalities are frequently incomplete. This highlights an urgent need to develop a robust
federated approach. However, we find that existing methods remain insufficient. On the one hand, centralized MGL methods that handle missing modalities overlook the knowledge sharing and generalization in
federated scenarios. On the other hand, while
federated MGL methods have become increasingly mature, they primarily target non-graph data. Based on these technologies, we identify a two-stage pipeline wherein client-side completion reconstructs missing modalities, and server-side aggregation integrates the client-updated parameters of both the modality generator and the backbone models. Although this serves as a general solution, we identify two primary challenges in achieving greater robustness: (1)
Topology-Isolated Local Completion: Client-side modality generation struggles to effectively leverage global semantics. (2) Reliability-Imbalanced Global Aggregation: Server-side multi-party collaboration is hindered by client updates with varying modality availability and recovery reliability. To address these challenges, we propose FedMPO, which utilizes
topology-aware cross-modal generation to recover missing features using comprehensive graph context, missing-aware expert routing to locally filter out noisy recovered signals, and reliability-aware aggregation to appropriately down-weight unreliable updates. Extensive experiments on 3 tasks across 6
datasets demonstrate that FedMPO outperforms baselines, achieving performance gains of up to 4.10% and 5.65% in high-missing and non-IID settings.
Read the full abstract ▾
All authors Sirui Zhang, Haonan Wang, Xunkai Li, Zekai Chen, Shumeng Li, Hongchao Qin, Rong-Hua Li, Guoren Wang
Machine Learning
CAWI: Copula-Aligned Weight Initialization for Randomized Neural Networks
Mushir Akhtar, M. Tanveer and Mohd. Arshad
Randomized neural networks (RdNNs) enable efficient, backpropagation-free training by freezing randomly initialized input-to-hidden weights, which permits a closed-form solution for the output layer. However, conventional random initialization is blind to inter-feature dependence, ignoring correlations, asymmetries, and tail dependence in the data, which degrades conditioning and predictive performance. To the best of our knowledge, this limitation remains unaddressed in the RdNN literature. To close this gap, we propose CAWI (Copula-Aligned Weight Initialization), a framework that draws input-to-hidden weights from a data-fitted copula that matches empirical dependence, ensuring the frozen projections respect inter-feature dependence without sacrificing the closed-form solution. CAWI (i) maps each feature to the unit interval using empirical CDFs, (ii) fits a multivariate copula that captures rank-based dependence among features, and (iii) samples each weight column w_j from the fitted copula and applies a fixed inverse marginal transform to set scale. The objective, solver, and "freeze-once" paradigm remain unchanged; only the sampling law for W becomes dependence-aware. For dependence modeling, we consider two copula families: elliptical (Gaussian, t) and Archimedean (Clayton, Frank, Gumbel). This enables CAWI to handle diverse dependence, including tail dependence. We evaluate CAWI across 83 diverse classification
benchmarks (binary and multiclass) and two biomedical
datasets, BreaKHis and the Schizophrenia
dataset, using standard shallow and deep RdNN architectures. CAWI consistently delivers significant improvements in predictive performance over conventional random initialization. Code is available at: https://github.com/mtanveer1/CAWI
Read the full abstract ▾
Machine Learning · Cross-listed from eess.SP
Recurrent Transformer-Based Near- and Far-Field THz Wideband Channel Estimation for UM-MIMO
Dmitry Artemasov, Alexander Shmatok, Kirill Andreev et al. (+3 more)
The integration of terahertz communications and ultra-massive multiple-input multiple-output (UM-
MIMO) systems in 6G networks is motivated by their ability to enable unprecedented data rates, mitigate spectrum congestion, and enhance overall network performance. However, the enlarged
antenna apertures and higher carrier frequencies in these systems increase the Rayleigh distance, causing users to span both the near-field and conventional far-field regions. Accurate spatial precoding thus requires exact channel estimation at the base station - a task made more challenging by the hybrid coexistence of near- and far-field effects and the limited number of digital chains available in hybrid beamforming architectures. In this paper, we propose a block recurrent
transformer model to address this challenge. We demonstrate that a single
transformer block equipped with state
memory can be trained once and then iteratively applied for hybrid-field channel estimation. Furthermore, we train the model such that it generalizes to wireless channels with varying scatterer distances, different numbers of propagation paths, and wideband operation. Simulation results show that the proposed method achieves performance gains of approximately 5 dB and 7.5 dB in normalized mean squared error (NMSE) over state-of-the-art solutions in narrowband and wideband scenarios, respectively.
Read the full abstract ▾
All authors Dmitry Artemasov, Alexander Shmatok, Kirill Andreev, Alexey Frolov, Manjesh K. Hanawal, Nikola Zlatanov
Machine Learning · Cross-listed from cs.CV
Improving Diffusion Posterior Samplers with Lagged Temporal Corrections for Image Restoration
Davide Evangelista, Elena Morotti, Francesco Pivi et al. (+1 more)
Diffusion-based posterior sampling (PS) is a leading framework for imaging inverse problems, combining learned priors with measurement constraints. Yet, its standard formulations rely on instantaneous data-consistent estimates, which induce temporal variability in the reverse dynamics. We reinterpret PS from a dynamical perspective, showing that the standard PS update corresponds to a first-order discretization of the
diffusion dynamics plus a residual correction capturing the mismatch between the denoised prediction and the data-consistent estimate. A second-order discretization, however, naturally introduces a temporal correction based on the variation of consecutive estimates. Building on this, we propose LAMP, combining the second-order update with the residual correction characterizing a PS technique. LAMP thus inherits a lagged temporal correction, and it can be implemented as a modular plug-in over the PS backbone. We show that LAMP preserves the structure of a posterior sampler, and we perform a one-step risk analysis to characterize when LAMP improves the reverse transition via a bias-variance trade-off. Experiments across multiple imaging tasks demonstrate consistent improvements over strong baselines such as DiffPIR and DDRM, without increasing the number of denoising evaluations.
Read the full abstract ▾
All authors Davide Evangelista, Elena Morotti, Francesco Pivi, Maurizio Gabbrielli
Machine Learning · Cross-listed from eess.IV
On Privacy-Preserving Image Transmission in Low-Altitude Networks: A Swin Transformer-Based Framework with Federated Learning
Kexin Zhang, Lixin Li, Yuna Yan et al. (+5 more)
The rapid development of low-altitude economy has driven the proliferation of Unmanned Aerial Vehicle (
UAV) applications, including logistics, inspection, and emergency response. However, transmitting high-volume image data from
UAVs to ground stations faces significant challenges due to limited bandwidth and stringent
privacy requirements. To address these issues, a Semantic Communication (SC) framework based on
Federated Learning (FL) is proposed for efficient and
privacy-preserving image transmission. A Swin
Transformer-based Semantic Communication (STSC) architecture is designed to extract multi-scale semantic features under constrained bandwidth conditions. Dedicated communication and computing nodes are deployed on
UAVs to enhance real-time coverage and flexibility. Meanwhile, a FL mechanism enables global model training across distributed devices without sharing raw data, thus preserving user
privacy. Simulation experiments conducted on the CIFAR-10
dataset demonstrate that the proposed STSC framework achieves at least 5.7 dB improvement in Peak Signal-to-Noise Ratio (PSNR) compared to DeepJSCC baselines, while also showing superior convergence and generalization performance. The framework effectively integrates
UAV-assisted deployment with SC and
privacy protection, offering a practical solution for bandwidth-constrained image transmission in low-altitude networks.
Read the full abstract ▾
All authors Kexin Zhang, Lixin Li, Yuna Yan, Xin Zhang, Wensheng Lin, Rui Li, Dongwei Zhao, Zhu Han
Machine Learning · Cross-listed from cs.CV
AlphaGRPO: Unlocking Self-Reflective Multimodal Generation in UMMs via Decompositional Verifiable Reward
Runhui Huang, Jie Wu, Rui Yang et al. (+2 more)
In this paper, we propose AlphaGRPO, a novel framework that applies Group Relative Policy Optimization (GRPO) to AR-
Diffusion Unified
Multimodal Models (UMMs) to enhance
multimodal generation capabilities without an additional cold-start stage. Our approach unlocks the model's intrinsic potential to perform advanced
reasoning tasks:
Reasoning Text-to-Image Generation, where the model actively infers implicit user intents, and Self-Reflective Refinement, where it autonomously diagnoses and corrects misalignments in generated outputs. To address the challenge of providing stable supervision for real-world
multimodal generation, we introduce the Decompositional Verifiable Reward (DVReward). Unlike holistic scalar rewards, DVReward utilizes an LLM to decompose complex user requests into atomic, verifiable semantic and quality questions, which are then evaluated by a general MLLM to provide reliable and interpretable feedback. Extensive experiments demonstrate that AlphaGRPO yields robust improvements across
multimodal generation
benchmarks, including GenEval, TIIF-Bench, DPG-Bench and WISE, while also achieving significant gains in editing tasks on GEdit without training on editing tasks. These results validate that our self-reflective reinforcement approach effectively leverages inherent understanding to guide high-fidelity generation. Project page: https://huangrh99.github.io/AlphaGRPO/
Read the full abstract ▾
All authors Runhui Huang, Jie Wu, Rui Yang, Zhe Liu, Hengshuang Zhao
Machine Learning
Pion: A Spectrum-Preserving Optimizer via Orthogonal Equivalence Transformation
Kexuan Shi, Hanxuan Li, Zeju Qiu et al. (+3 more)
We introduce Pion, a spectrum-preserving optimizer for large
language model (LLM) training based on orthogonal equivalence transformation. Unlike additive optimizers such as Adam and
Muon, Pion updates each weight
matrix through left and right orthogonal transformations, preserving its singular values throughout training. This yields an optimization mechanism that modulates the geometry of weight
matrices while keeping their spectral norm fixed. We derive the Pion update rule, systematically examine its design choices, and analyze its convergence behavior along with several key properties. Empirical results show that Pion offers a stable and competitive alternative to standard optimizers for both LLM pretraining and finetuning.
Read the full abstract ▾
All authors Kexuan Shi, Hanxuan Li, Zeju Qiu, Yandong Wen, Simon Buchholz, Weiyang Liu
Machine Learning · Cross-listed from cs.CV
Do Vision Transformers Need All-to-All Attention? Global Communication Through Elastic Learned Cores
Alan Z. Song, Yinjie Chen, Mu Nan et al. (+3 more)
Vision Transformers (ViTs) learn rich visual-semantic representations through all-to-all self-
attention among patch
tokens. However, this design implicitly assumes that direct pairwise patch interactions are necessary for effective representation learning. In this work, we challenge this assumption and show that representations supporting both global recognition and dense prediction can be learned without direct patch-to-patch interaction. We propose VECA (Visual Elastic-Core
Attention), a
vision transformer with core-periphery structured
attention mediated by a small set of learned cores. Patch
tokens exchange global information exclusively through these cores, while the full set of dense patches are preserved and iteratively updated across layers. This reduces
attention complexity from $O(N^2)$ to $O(N)$, linear in the number of patches $N$ for a fixed core budget $C$. Unlike prior latent-
token cross-
attention architectures, VECA facilitates sparse global communication without compressing the spatial representation itself. Nested training along the core axis further enables a single model to elastically trade off computation and accuracy at inference time without retraining. Across image classification and dense prediction tasks, VECA remains competitive with full-
attention backbones and outperforms the evaluated linear-complexity alternatives on most
benchmarks. Moreover, without explicit supervision, these cores develop semantically organized structures that support object-label transfer across
video frames. These results show that effective visual representations can be learned without direct all-to-all patch interaction.
Read the full abstract ▾
All authors Alan Z. Song, Yinjie Chen, Mu Nan, Deva Ramanan, Michael J. Tarr, Andrew F. Luo
Machine Learning · Cross-listed from cs.CL
Task-Adaptive Embedding Refinement via Test-time LLM Guidance
Ariel Gera, Shir Ashury-Tahan, Gal Bloch et al. (+2 more)
We explore the effectiveness of an LLM-guided query refinement paradigm for extending the usability of
embedding models to challenging zero-shot search and classification tasks. Our approach refines the
embedding representation of a user query using feedback from a generative LLM on a small set of documents, enabling
embeddings to adapt in real time to the target task. We conduct extensive experiments with state-of-the-art text
embedding models across a diverse set of challenging search and classification
benchmarks. Empirical results indicate that LLM-guided query refinement yields consistent gains across all models and
datasets, with relative improvements of up to +25% in literature search, intent detection, key-point matching, and nuanced query-instruction following. The refined queries improve ranking quality and induce clearer binary separation across the corpus, enabling the
embedding space to better reflect the nuanced, task-specific constraints of each ad-hoc user query. Importantly, this expands the range of practical settings in which
embedding models can be effectively deployed, making them a compelling alternative when costly LLM pipelines are not viable at corpus-scale. We release our experimental code for reproducibility, at https://github.com/IBM/task-aware-
embedding-refinement.
Read the full abstract ▾
All authors Ariel Gera, Shir Ashury-Tahan, Gal Bloch, Ohad Eytan, Assaf Toledo
Machine Learning
Learning, Fast and Slow: Towards LLMs That Adapt Continually
Rishabh Tiwari, Kusha Sareen, Lakshya A Agrawal et al. (+6 more)
Large
language models (LLMs) are trained for downstream tasks by updating their parameters (e.g., via RL). However, updating model parameters forces them to absorb task-specific information, which can result in catastrophic forgetting and loss of plasticity. In contrast, in-context learning with fixed LLM parameters can cheaply and rapidly adapt to task-specific requirements (e.g.,
prompt optimization), but cannot by itself typically match the performance gains available through updating LLM parameters. Inspired by how human learning works at multiple time scales (e.g., System 1 vs 2), we propose a fast-slow learning framework for LLMs, instantiated with model parameters as "slow" weights and optimized context as "fast" weights. These fast "weights" can learn from textual feedback to absorb the task-specific information, while allowing slow weights to stay closer to the base model and persist general
reasoning behaviors. Fast-Slow Training (FST) is up to 3x more sample-efficient than only slow learning (RL) across
reasoning tasks, while consistently reaching a higher performance asymptote. Moreover, FST-trained models remain closer to the base LLM (up to 70% less KL divergence), resulting in less catastrophic forgetting than RL-training. This reduced drift also preserves plasticity: after training on one task, FST trained models adapt more effectively to a subsequent task. In continual learning scenarios, where task domains change on the fly, FST continues to acquire each new task while RL stalls.
Read the full abstract ▾
All authors Rishabh Tiwari, Kusha Sareen, Lakshya A Agrawal, Joseph E. Gonzalez, Matei Zaharia, Kurt Keutzer, Inderjit S Dhillon, Rishabh Agarwal, Devvrit Khatri