Machine Learning
Sparse Prefix Caching for Hybrid and Recurrent LLM Serving
Mikhail Shirokikh and Sergey Nikolenko
Prefix caching is a key latency optimization for autoregressive LLM serving, yet existing systems assume dense per-
token key/value reuse. State-space models change the structure of the problem: a recurrent layer can resume from a single stored state rather than requiring the entire
token history. This asymmetry opens a new design point between no reuse and dense caching: store exact recurrent states at a sparse set of checkpoint positions and, on a cache hit, resume from the deepest stored checkpoint and recompute the remaining suffix exactly. We formalize sparse prefix caching as checkpoint placement under a distribution over overlap depths, yielding an exact O(NM) dynamic program. For use cases where requests share a non-trivial prefix (e.g. asking different questions about a single long document), we show that our method consistently improves the Pareto frontier traced by standard heuristics on real-world data. Across QuALITY and System
Prompts, distribution-aware placement dominates every fixed-budget baseline on the measured layer-group Pareto frontier and matches or outperforms the strongest heuristic (block caching) while typically using substantially fewer checkpoints, with the largest gains at low checkpoint budgets where the overlap distribution is most non-uniform. The method is most relevant when many requests share a substantial but not identical prefix within a retained cache entry. It preserves exact outputs, does not change the recurrent computation itself or require new recurrent update kernels, applies to recurrent/SSM layers whose hidden state can be extracted and restored exactly, and for hybrid models can be combined with existing KV-cache compression techniques.
Read the full abstract ▾
Machine Learning
Horizon-Constrained Rashomon Sets for Chaotic Forecasting
Gauri Kale, Rahul Vishwakarma, Holly Diamond et al. (+2 more)
Predictive multiplicity and chaotic dynamics represent two fundamental challenges in machine learning that have evolved independently despite their conceptual connections. We bridge this gap by introducing horizon-constrained Rashomon sets, a theoretical framework that characterizes how model multiplicity evolves with prediction horizon in chaotic systems. Unlike static prediction tasks where the Rashomon set remains fixed, chaos induces exponential divergence among initially similar models, fundamentally transforming the nature of predictive equivalence. We prove that the effective Rashomon set contracts exponentially with lead time at a rate determined by the maximum Lyapunov exponent and introduce Lyapunov-weighted metrics that provide tighter bounds on predictive disagreement. Leveraging these insights, we develop decision-aligned selection
algorithms that choose among near-optimal models based on downstream utility rather than forecast accuracy alone. Extensive experiments on synthetic chaotic systems (Lorenz-96, Kuramoto-Sivashinsky) and real-world applications (wind power, traffic, weather) demonstrate that our framework improves decision quality by 18-34% while maintaining competitive predictive performance. This work establishes the first rigorous connection between chaos theory and predictive multiplicity, providing principled guidance for deploying machine learning in safety-critical chaotic domains.
Read the full abstract ▾
All authors Gauri Kale, Rahul Vishwakarma, Holly Diamond, Ava Hedayatipour, Amin Rezaei
Machine Learning
Physics-Informed Neural Networks with Learnable Loss Balancing and Transfer Learning
Reza Pirayeshshirazinezhad
We propose a self-supervised physics-informed neural network (PINN) framework that adaptively balances physics-based and data-driven supervision for scientific machine learning under data scarcity. Unlike prior PINNs that rely on fixed or heuristic weighting of physics residuals and data loss, our approach introduces a learnable blending
neuron that dynamically adjusts the relative contribution of each term based on their uncertainties. This mechanism enables stable training and improved generalization without manual tuning. To further enhance efficiency, we integrate a transfer learning strategy that reuses representations from related domains and adapts them to new physical systems with limited data. We validate the framework for the prediction of heat transfer in
liquid-metal miniature heat sinks using only 87 CFD datapoints, where the adaptive PINN achieves an error <8%, outperforming shallow neural networks, kernel methods, and physics-only baselines. Our framework provides a general recipe for
embedding physics adaptively into neural networks, offering a robust and reproducible approach for data-scarce problems across various scientific domains, including fluid dynamics and material modeling.
Read the full abstract ▾
Machine Learning
SAT: Sequential Agent Tuning for Coordinator Free Plug and Play Multi-LLM Training with Monotonic Improvement Guarantees
Yi Xie, Yangyang Xu, Yi Fan et al. (+1 more)
Large
language models (LLMs) with a large number of parameters achieve strong performance but are often prohibitively expensive to deploy. Recent work explores using teams of smaller, more efficient LLMs that collectively match or even outperform a single large model. However, jointly updating multiple
agents introduces compounding distribution shifts, making coordination and stability during training difficult. We address this by introducing Sequential
Agent Tuning (SAT), a coordinator-free training paradigm. SAT represents the team as a factorized policy and employs block-coordinate updates over
agents, enabling scalable, decentralized training without a central controller. Specifically, we develop a sequence-aware, on-policy advantage
estimator that conditions on the evolving team policy, coupled with per-
agent KL trust regions that isolate occupancy drift. Theoretically, this framework provides two critical guarantees. First, it ensures monotonic improvement, stabilizing the training process. Second, it establishes provable plug-and-play invariance: any
agent can be upgraded to a stronger model without retraining the rest of the team, with a formal guarantee that the performance bound improves. Empirically, a team of three 4B
agents (12B total) trained with SAT surpasses the much larger Qwen3-32B on AIME24/25
benchmarks by 3.9% on average. We validate our plug-and-play theory by swapping in two 8B
agents, which boosts the composite score by 10.4%. We provide code and appendix of proof at https://github.com/Yydc/SAT-AAMAS
Read the full abstract ▾
All authors Yi Xie, Yangyang Xu, Yi Fan, Bo Liu
Machine Learning · Cross-listed from cs.CV
Layout-Aware Representation Learning for Open-Set ID Fraud Discovery
Jinxing Li, Nicholas Ren, Cathy Chang et al. (+2 more)
Identity-document fraud detection is not a stationary binary classification problem. Adaptive attackers modify templates and fabrication pipelines, making historical fraud labels stale, and successful forgeries recur at scale as coherent campaigns. We therefore study layout-aware representation learning for open-set fraud discovery rather than only closed-set classification. We adapt DINOv3 to the document domain via context-aware SimMIM fine-tuning and supervised metric learning with composite loss that encourages inter-class separability and intra-class compactness. The model is trained with U.S. IDs only. With a lightweight MLP and softmax classifier, the
embedding achieves 99.83% layout classification accuracy on Canadian layouts. Moreover, on a
dataset of 20,448 Canadian IDs,
embedding-space analysis surfaces 276 adaptive physical-fraud cases, including 222 not surfaced by incumbent detectors. The
embedding supports similarity-based expansion from a single confirmed seed to additional related cases not linked by conventional metadata graphs. The layout-aware document
embeddings provide a production-aligned basis for discovering novel and campaign-scale fraud under distribution shift.
Read the full abstract ▾
All authors Jinxing Li, Nicholas Ren, Cathy Chang, Hongkai Pan, Daniel George
Machine Learning · Cross-listed from eess.SP
MedMamba: Recasting Mamba for Medical Time Series Classification
ZhengXiao He, Huayu Li, Xiwen Chen et al. (+4 more)
Medical time series, such as electrocardiograms (ECG) and electroencephalograms (
EEG), exhibit complex temporal dynamics and structured cross-channel dependencies, posing fundamental challenges for automated analysis. Conventional convolutional and recurrent models struggle to capture long-range dependencies, while
Transformer-based approaches incur quadratic complexity and often introduce redundant interactions that are misaligned with the intrinsic structure of physiological signals. To address these limitations, we propose MedMamba, a principle-driven multi-scale bidirectional state space architecture tailored for medical time series classification. Our design is guided by three key inductive biases of physiological signals: spatial centralization, multi-timescale temporal composition, and non-causal contextual dependency. These principles are instantiated through a lightweight channel-mixing module for cross-channel reparameterization, multi-scale convolutional tokenization for temporal decomposition, and bidirectional Mamba blocks for efficient global context modeling with linear complexity. Extensive experiments on six
benchmark datasets spanning
EEG, ECG, and human activity signals demonstrate that MedMamba consistently outperforms state-of-the-art methods across diverse modalities. Notably, it achieves 85.97% accuracy on PTB and establishes new state-of-the-art performance on the challenging ADFTD
dataset (54.72% accuracy and 52.01% F1-score). Strong results on long-sequence
benchmarks, such as SleepEDF, further validate its capability in modeling long-range dependencies. Moreover, MedMamba achieves a speedup of 4.6x in inference, highlighting its practicality for real-time clinical deployment. These results suggest that principle-guided state space modeling offers an effective and scalable alternative to
Transformer-based approaches for medical time series analysis.
Read the full abstract ▾
All authors ZhengXiao He, Huayu Li, Xiwen Chen, Janet M Roveda, Jinghao Wen, Siyuan Tian, Ao Li
Machine Learning
A Randomized PDE Energy driven Iterative Framework for Efficient and Stable PDE Solutions
Yi Bing, Zheng Ran, Fu Jinyang et al. (+2 more)
Efficient and stable solution of partial differential equations (PDEs) is central to scientific and engineering applications, yet existing numerical solvers rely heavily on
matrix based discretizations, while learning based methods require costly training and often suffer from limited generalization. In this work, we proposes a PDE energy driven framework that solves PDEs through physically constrained
diffusion iterations, without relying on classical
matrix based finite element assembly or data driven neural network training. The proposed method evolves arbitrary random initial fields through PDE energy driven implicit iterations combined with Gaussian smoothing, while strictly enforcing boundary conditions at each iteration. The proposed formulation is applied to representative one dimensional Poisson, Heat, and viscous Burgers equations, covering both steady state and transient problems. Numerical results demonstrate stable convergence to the unique physical solution from random initializations, with accurate resolution of sharp gradients and controlled Mean Squared Error (MSE) across a wide range of discretization parameters. Detailed comparisons with analytical solutions indicate that the framework achieves competitive accuracy and stability. Overall, the proposed framework provides a fast, flexible, and physically consistent alternative to traditional numerical solvers, offering a potential pathway for scalable PDE solutions in both research and engineering applications.
Read the full abstract ▾
All authors Yi Bing, Zheng Ran, Fu Jinyang, Liu Long, Peng Xiang
Machine Learning
A Multimodal and Explainable Machine Learning Approach to Diagnosing Multi-Class Ejection Fraction from Electrocardiograms
Catherine Ning, Yu Ma, Cindy Beini Wang et al. (+4 more)
Left ventricular ejection fraction (LVEF) assessment depends on echocardiography, limiting access in primary care and resource-constrained settings. We developed a
multimodal machine-learning framework that combines engineered 12-lead ECG timeseries features with structured EHR variables to classify LVEF into four clinically used strata: normal (>50%), mildly reduced (40-50%), moderately reduced (30-40%), and severely reduced (<30%). To support model explainability, we identified the most influential ECG and EHR features via SHAP attributions. Using retrospective data from Hartford HealthCare, we trained XGBoost models on 36,784 ECG-echocardiogram pairs from 30,952 outpatients and evaluated temporal generalizability on 19,966 ECGs from a subsequent period. The
multimodal model achieved one-vs-rest AUROCs of 0.95 (severe), 0.92 (moderate), 0.82 (mild), and 0.91 (normal), outperforming ECG-only and EHR-only baselines, and maintained performance under temporal validation. This work supports ECG-based,
multimodal LVEF stratification as a practical screening and triage aid to prioritize confirmatory imaging where resources are limited.
Read the full abstract ▾
All authors Catherine Ning, Yu Ma, Cindy Beini Wang, Sean McMahon, Joseph Radojevic, Steven Zweibel, Dimitris Bertsimas
Machine Learning · Cross-listed from physics.data-an
Application of a Mixture of Experts-based Foundation Model to the GlueX DIRC Detector
Cristiano Fanelli, James Giroux, Cole Granger et al. (+1 more)
We present a Mixture-of-Experts-based foundation model applied to the GlueX DIRC detector at Jefferson Lab, demonstrating its utility as a unified framework for fast simulation, particle identification, and hit-level noise filtering of Cherenkov
photons. By leveraging a single shared
transformer backbone across all tasks, the approach eliminates the fragmentation of task-specific pipelines while maintaining competitive-and in several cases superior-performance relative to established methods. The model operates directly on low-level detector inputs, performing hit-by-hit autoregressive generation over split spatial and temporal vocabularies with continuous kinematic conditioning, and supports class-conditional generation of pions and kaons through its Mixture-of-Experts architecture. We
benchmark against the standard geometrical reconstruction and prior deep learning methods across the full kinematic phase space of the GlueX DIRC, demonstrating that the foundation model framework transfers effectively to this detector without architectural modification. This work positions the foundation model as a practical and scalable alternative to the suite of task-specific models currently proposed for GlueX DIRC analysis.
Read the full abstract ▾
All authors Cristiano Fanelli, James Giroux, Cole Granger, Justin Stevens
Machine Learning · Cross-listed from cs.SD
Audio2Tool: Speak, Call, Act – A Dataset for Benchmarking Speech Tool Use
Ramit Pahwa, Apoorva Beedu, Parivesh Priye et al. (+4 more)
Voice assistants increasingly rely on
Speech Language Models (SpeechLMs) to interpret spoken queries and execute complex tasks, yet existing
benchmarks lack domain breadth, acoustic diversity, and compositional
reasoning complexity to evaluate tool-calling performance. We introduce Audio2Tool, a large-scale
dataset comprising approximately 30,000 queries designed to assess tool-calling capabilities of SpeechLMs across three primary domains: Smart Car, Smart Home, and Wearables. Our
benchmark features a multi-tier complexity hierarchy, ranging from simple direct commands to complex multi-intent and needle-in-a-haystack extraction to isolate distinct failure modes. To ensure realism, we employ zero-shot voice cloning text-to-
speech synthesis and diverse noise profiles to simulate in-the-wild conditions. Evaluations of state-of-the-art SpeechLMs and ASR-LLM pipelines show strong performance on simple commands but significant degradation under compositional and acoustic challenges. Code and
dataset are publicly available on the project page: https://audio2tool.github.io/.
Read the full abstract ▾
All authors Ramit Pahwa, Apoorva Beedu, Parivesh Priye, Rutu Gandhi, Saloni Takawale, Aruna Baijal, Zengli Yang
Machine Learning · Cross-listed from math.DS
Spectral Kernel Dynamics for Planetary Surface Graphs: Distinction Dynamics and Topological Conservation
Jnaneshwar Das
The spectral kernel field equation R[k] = T[k] lacks a conservation-law analog. We prove (i) the fixed-point flow is strictly volume-expanding (tr DF > 0), precluding automatic conservation, and (ii) the conservation deficit per mode equals the Hessian stability margin exactly: D_m = -Delta'. Closing the deficit requires a scene-side compensating contribution, which we formalise as the distinction dynamics equation dc/dt = G[c, h_t], with MaxCal-optimal realisation G_opt. On fixed-
topology 3D surface graphs we derive a conditional
topology-preserving compression
theorem: retaining k >= beta_0 + beta_1 modes (under a spectral-ordering assumption) preserves all Betti-number charges; we include a worked short-cycle counterexample (figure-eight) calibrating when the assumption fails. A triple necessary spectral diagnostic – Fiedler-mode concentration, elevated curl energy, anomalous beta_1 – is derived for planetary drainage networks at O(N) cost. Two internal real-data sequences serve as preliminary consistency checks; full
benchmarks and adaptive-
topology extensions are deferred.
Read the full abstract ▾
Machine Learning · Cross-listed from physics.chem-ph
KinetiDiff: Docking-Guided Diffusion for De Novo ACVR1 Inhibitor Design in Fibrodysplasia Ossificans Progressiva
Aaryan Patel
We present KinetiDiff, a structure-based framework for de novo kinase inhibitor design that integrates a Geometry-Complete
Diffusion Model with real-time AutoDock Vina gradient guidance. By injecting physics-based docking gradients into the
diffusion denoising loop, KinetiDiff steers molecule generation toward high-affinity conformations for ACVR1 (ALK2), the causative kinase in Fibrodysplasia Ossificans Progressiva. From 10,000
diffusion samples, the framework produced 9,997 valid molecules. The best candidate achieved $-11.05$ kcal/mol (pKd = 8.10), a 19.2% improvement over the crystallographic reference. The top 100 candidates all exceed the reference, with 100% Lipinski compliance, median synthetic accessibility of 2.67, and internal diversity of 0.790. Systematic ablation across four guidance strategies–Vina-Direct (physics), HNN-Denovo (neural proxy), multi-objective, and unguided–demonstrates that real-time docking guidance dominates on all metrics. We evaluate HNN-Denovo as a computationally efficient alternative (60-fold speedup per step), revealing a domain-mismatch limitation (r = 0.224 correlation with Vina) that explains its inferior performance. These results establish gradient-guided
geometric diffusion as a practical approach for generating potent, synthetically accessible inhibitors against rare-disease kinase targets.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.ET
What Makes a Bacterial Model a Good Reservoir Computer? Predicting Performance from Separability and Similarity
Laura Alonso Bartolomé, Jean-Loup Faulon and Xavier Hinaut
Biological systems are promising substrates for computation because they naturally process environmental information through complex internal dynamics. In this study, we investigate whether bacterial metabolic models can act as physical reservoirs and whether their computational performance can be predicted from dynamical properties linked to separability and similarity. We simulated the growth dynamics of five bacterial species, one yeast species, and 29 Escherichia coli single-gene deletion mutants using dynamic flux balance analysis (dFBA), with glucose and xylose concentrations as inputs and growth curves as reservoir states. Computational performance was assessed on random nonlinear classification tasks using a linear readout, while reservoir properties linked to separability and similarity were characterised through kernel and generalisation ranks computed from growth-curve state
matrices. Several microbial models achieved high classification accuracy, showing that bacterial metabolic dynamics can support nonlinear computation. Clear differences were observed between species, with some models converging more rapidly and others reaching higher maximum accuracy, revealing a trade-off between convergence speed and peak performance. In contrast, all E. coli mutants were dominated by the wild-type model, suggesting that gene deletions reduce the dynamical richness required for efficient computation. The difference between kernel and generalisation ranks was generally associated with improved accuracy, but deviations across models and sensitivity at low rank values limited its predictive power in practice. Overall, these results show that bacterial metabolic models constitute promising substrates for reservoir computing and provide a first step towards identifying microbial strains with favourable computational properties for future experimental implementations.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.AI
Debate as Reward: A Multi-Agent Reward System for Scientific Ideation via RL Post-Training
Moein Salimi, Babak Hosseini Mohtasham, Amin Aghakasiri et al. (+6 more)
Large
Language Models (LLMs) have demonstrated potential in automating scientific ideation, yet current approaches relying on iterative prompting or complex multi-
agent architectures often suffer from
hallucination or computational inefficiency. A critical bottleneck in applying Reinforcement Learning (RL) to this open-ended domain is reward hacking – where models exploit imperfect evaluation proxies to maximize scores without producing genuine scientific innovation. To address these limitations, we propose an RL framework explicitly tailored for high-quality scientific idea generation. We propose the first multi-
agent reward function designed to serve as a judge, decoupling methodological validation from implementation details while providing strict binary rewards that are robust to reward hacking. To effectively optimize against this sparse signal, we utilize an unbiased variant of Group Relative Policy Optimization to mitigate artificial length bias. We grounded our training in ICLR-320, a curated
dataset of problem-solution pairs extracted from ICLR 2024 proceedings. Experiments demonstrate that our framework significantly outperforms state-of-the-art baselines across expert-evaluated metrics of novelty, feasibility, and effectiveness.
Read the full abstract ▾
All authors Moein Salimi, Babak Hosseini Mohtasham, Amin Aghakasiri, Mahdi Naieni, Amir Hossein Qeysarbeigi, Mohammad Masih Shalchian Nazer, Zahra Azar, Mahdi Jafari Siavoshani, Mohammad Hossein Rohban
Machine Learning
Neuroscience Inspired Graph Operators Towards Edge-Deployable Virtual Sensing for Irregular Geometries
William Howes, Farid Ahmed, Kazuma Kobayashi et al. (+2 more)
Predicting full-field physics through the real-time virtual sensing of engineering systems can enhance limited physical sensors but often requires sparse-to-dense reconstruction, complex multiphysics, and highly irregular geometries as well as strict latency and energy constraints for edge-deployability. Neural operators have been presented as a potential candidate for such applications but few architectures exist that explicitly address power consumption. Spiking
neuron integration can provide a potential solution when integrated on neuromorphic hardware but the current existing
neuron models result in severe performance degradation towards regression-based virtual sensing. To address the performance concerns and edge-constraints, we present the Variable Spiking Graph Neural Operator (VS-GNO) which integrates a sophisticated spectral-spatial convolutional analysis and a previously developed Variable Spiking
Neuron (VSN) and energy-error balance loss function. With a non-spiking $L_2$ error baseline of $0.4\%$, VS-GNO can provide a reconstruction error of $0.71\%$ with $15\%$ average spiking in its spectral-only form and $1.04\%$ with $24.5\%$ spiking in its entire form. These results position VS-GNO as a promising step towards energy-efficient, edge-deployable neural operators for real-time sparse-to-dense virtual sensing in complex, highly irregular engineering environments.
Read the full abstract ▾
All authors William Howes, Farid Ahmed, Kazuma Kobayashi, Souvik Chakraborty, Syed Bahauddin Alam
Machine Learning
Late Fusion Neural Operators for Extrapolation Across Parameter Space in Partial Differential Equations
Eva van Tegelen, Taniya Kapoor, George A.K. van Voorn et al. (+2 more)
Developing neural operators that accurately predict the behavior of systems governed by partial differential equations (PDEs) across unseen parameter regimes is crucial for robust generalization in scientific and engineering applications. In practical applications, variations in physical parameters induce distribution shifts between training and prediction regimes, making extrapolation a central challenge. As a result, the way parameters are incorporated into neural operator models plays a key role in their ability to generalize, particularly when state and parameter representations are entangled. In this work, we introduce the Late Fusion Neural Operator, an architecture that disentangles learning state dynamics from parameter effects, improving predictive performance both within and beyond the training distribution. Our approach combines neural operators for learning latent state representations with sparse regression to incorporate parameter information in a structured manner. Across four
benchmark PDEs including advection, Burgers, and both 1D and 2D reaction-
diffusion equations, the proposed method consistently outperforms
Fourier Neural Operator and CAPE-FNO. Late Fusion Neural Operators achieve consistently the best performance in all experiments, with an average RMSE reduction of 72.9% in-domain and 71.8% out-domain compared to the second-best method. These results demonstrate strong generalization across both in-domain and out-domain parameter regimes.
Read the full abstract ▾
All authors Eva van Tegelen, Taniya Kapoor, George A.K. van Voorn, Peter van Heijster, Ioannis N. Athanasiadis
Machine Learning
Chronax: A Jax Library for Forecasting and Conformal Inference
Xan Carey, Arsh Singh, Yash Deshmukh et al. (+8 more)
Time-series
forecasting is central to many scientific and industrial domains, such as energy systems,
climate modeling, finance, and retail. While
forecasting methods have evolved from classical statistical models to automated, and neural approaches, the surrounding software ecosystem remains anchored to the traditional Python numerical stack. Existing libraries rely on interpreter-driven execution and object-oriented abstractions, limiting composability, large-scale parallelism, and integration with modern differentiable and accelerator-oriented workflows. Meanwhile, today's
forecasting increasingly involves large collections of heterogeneous time series data, irregular
covariates, and frequent retraining, placing new demands on scalability and execution efficiency. JAX offers an alternative paradigm to traditional stateful numerical computation frameworks based on pure functions and program transformations such as just-in-time compilation and automatic vectorization, enabling end-to-end optimization across CPUs, GPUs, and TPUs. However, this modern paradigm has not yet been fully incorporated into the design of
forecasting systems. We introduce Chronax, a JAX-native time-series
forecasting library that rethinks
forecasting abstractions around functional purity, composable transformations, and accelerator-ready execution. By representing preprocessing, modeling, and multi-horizon prediction as pure JAX functions, Chronax enables scalable multi-series
forecasting, model-agnostic conformal uncertainty quantification, and seamless integration with modern machine learning and scientific computing pipelines.
Read the full abstract ▾
All authors Xan Carey, Arsh Singh, Yash Deshmukh, Sunit Jadhav, Omkar Tekawade, Aileen Huang, Lorraine Yang, Anvesha Tiwary, Gerardo Riano, Amy Greenwald, Denizalp Goktas
Machine Learning · Cross-listed from cs.DC
Scalable and Adaptive Parallel Training of Graph Transformer on Large Graphs
Jun-Liang Lin, Kamesh Madduri and Mahmut Taylan Kandemir
Graph foundation models have demonstrated remarkable adaptability across diverse downstream tasks through large-scale pretraining on graphs. However, existing implementations of the backbone model, graph
transformers, are typically limited to single-GPU systems, leading to long training times or out-of-
memory issues on large graphs. Moreover, parallelizing graph
transformer training over the full graph is challenging, as efficiency depends heavily on both the graph structure and system characteristics, such as bandwidth and
memory capacity. In this work, we introduce a distributed training framework for graph
transformers, which automatically selects and optimizes parallelization strategies based on the graph structure and hardware configuration. With our implementation of distributed sparse operations, we accelerate sparse graph
attention by up to 3.8x and reduce
memory consumption by 78% compared to state-of-the-art frameworks. On large graph
benchmarks, our proposed framework achieves up to 6x speedup with system scaling up to 8 GPUs. These results demonstrate that the proposed framework improves the scalability of graph
transformers, bringing them closer to serving as practical graph foundation models.
Read the full abstract ▾
Machine Learning
How to Approximate Inference with Subtractive Mixture Models
Lena Zellinger, Nicola Branchini, Lennert De Smet et al. (+3 more)
Classical mixture models (MMs) are widely used tractable proposals for approximate inference settings such as variational inference (VI) and importance sampling (IS). Recently, mixture models with negative coefficients, called subtractive mixture models (SMMs), have been proposed as a potentially more expressive alternative. However, how to effectively use SMMs for VI and IS is still an open question as they do not provide latent variable semantics and therefore cannot use sampling schemes for classical MMs. In this work, we study how to circumvent this issue by designing several expectation
estimators for IS and learning schemes for VI with SMMs, and we empirically evaluate them for distribution approximation. Finally, we discuss the additional challenges in estimation stability and learning efficiency that they carry and propose ways to overcome them. Code is available at: https://github.com/april-tools/delta-vi.
Read the full abstract ▾
All authors Lena Zellinger, Nicola Branchini, Lennert De Smet, Víctor Elvira, Nikolay Malkin, Antonio Vergari
Machine Learning · Cross-listed from cs.CR
Surgical Repair of Insecure Code Generation in LLMs
Gustavo Sandoval, Brendan Dolan-Gavitt and Siddharth Garg
Large
language models write production code, and yet they routinely introduce well-known vulnerabilities. We show that this is not a knowledge deficit: the same models that generate insecure code, correctly identify and explain the vulnerability when asked directly, this is a gap we call the Format-Reliability Gap. Mechanistic analysis reveals the cause: security representations are encoded from the earliest layers but remain computationally inert until the final layer, where format-compliance demands compete with them. Because the failure is localized to a single layer, per-vulnerability steering vectors reduce insecure generation by up to 74% with negligible overhead. The mechanism and the fix generalize across five models, three architecture families, and six vulnerability types, suggesting insecure code generation is an interpretability problem, not a training artifact.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.AI
Agentic Risk-Aware Set-Based Engineering Design
Varun Kumar and George Em Karniadakis
This paper introduces a multi-
agent framework guided by Large
Language Models (LLMs) to assist in the early stages of engineering design, a phase often characterized by vast parameter spaces and inherent uncertainty. Operating under a human-in-the-loop paradigm and demonstrated on the canonical problem of aerodynamic airfoil design, the framework employs a team of specialized
agents: a Coding Assistant, a Design
Agent, a Systems Engineering
Agent, and an Analyst
Agent - all coordinated by a human Manager. Integrated within a set-based design philosophy, the process begins with a collaborative phase where the Manager and Coding Assistant develop a suite of validated tools, after which the
agents execute a structured workflow to systematically explore and prune a large set of initial design candidates. A key contribution of this work is the explicit integration of formal risk management, employing the Conditional Value-at-Risk (CVaR) as a quantitative metric to filter designs that exhibit a high probability of failing to meet performance requirements, specifically the target coefficient of lift. The framework automates labor-intensive initial exploration through a global sensitivity analysis conducted by the Analyst
agent, which generates actionable heuristics to guide the other
agents. The process culminates by presenting the human Manager with a curated final set of promising design candidates, augmented with high-fidelity Computational Fluid Dynamics (CFD) simulations. This approach effectively leverages AI to handle high-volume analytical tasks, thereby enhancing the decision-making capability of the human expert in selecting the final, risk-assessed design.
Read the full abstract ▾
Machine Learning
Graph Transformer-Based Pathway Embedding for Cancer Prognosis
Koushik Howlader, Md Tauhidul Islam and Wei Le
Accurate prediction of cancer progression remains a challenge due to the high heterogeneity of molecular omics data across patients. While biologically informed models have improved the interpretability of these predictions, a persistent limitation lies in how they encode individual genes to construct pathway representations. Existing hierarchical models typically derive gene features by directly mapping raw molecular inputs, whereas integration frameworks often rely on simple statistical aggregations of patient-level signals. These approaches often fail to explicitly learn a shared base representation for each gene, thereby limiting the expressiveness and biological accuracy of downstream pathway
embeddings. To address this, we introduce PATH, a modulation-based, patient-conditioned gene
embedding strategy. PATH represents a paradigm shift by starting from a shared base
embedding for each gene, preserving a stable biological identity across the population, and then dynamically adapting it using patient-specific copy number variation (CNV) and mutation signals. This allows the model to capture subtle individual molecular variations while maintaining a consistent latent understanding of the gene itself. We integrate PATH into a graph
transformer framework that models interactions among biologically connected pathways through pathway-guided
attention. Across pancancer metastasis prediction, PATH achieves an F1 score of 0.8766, representing an 8.8 percent improvement over the current SOTA multi-omics
benchmarks. Beyond superior predictive accuracy, our approach identifies biologically meaningful pathways and, crucially, reveals disease-state-specific pathway rewiring, offering new insights into the evolving pathway-pathway interactions that drive cancer progression.
Read the full abstract ▾
Machine Learning
DARLING: Detection Augmented Reinforcement Learning with Non-Stationary Guarantees
Argyrios Gerogiannis, Yu-Han Huang and Venugopal V. Veeravalli
We study model-free reinforcement learning (RL) in non-stationary finite-horizon episodic Markov decision processes (MDPs) without prior knowledge of the non-stationarity. We focus on the piecewise stationary (PS) setting, where both rewards and transition dynamics can change at unknown times. We first revisit existing state-of-the-art approaches and identify theoretical and practical limitations that change the current landscape of performance guarantees. To characterize the difficulty of the problem, we establish the first minimax lower bounds for PS-RL in tabular and linear MDPs. We then introduce Detection Augmented Reinforcement Learning (DARLING), a modular wrapper for PS-RL that applies to both tabular and linear MDPs, without knowledge of the changes. In tabular MDPs, under change-point separability and reachability conditions, DARLING improves the best known dynamic regret bounds and matches our minimax lower bound. In linear MDPs, DARLING matches the minimax lower bound when the relevant reachability parameters are known, and our analysis clarifies the structural obstacles that distinguish this setting from the tabular case. Finally, through extensive experimentation across diverse non-stationary
benchmarks, we show that DARLING consistently surpasses the state-of-the-art methods.
Read the full abstract ▾
Machine Learning
UniCon: Unified Framework for Efficient Contrastive Alignment via Kernels
Hangke Sui, Yuqing Wang and Minh N Do
Contrastive objectives power state-of-the-art
multimodal models, but their training remains slow, relying on long
stochastic optimization. We propose a Unified Framework for Efficient Contrastive
Alignment via Kernels (UniCon), which spans linear and nonlinear encoders as well as one-to-one and many-to-many
alignments. At its core, UniCon introduces the contrastive similarity weight
matrix $S(\gamma)$, which enables closed-form global solutions that provably replace minibatch back-propagation with exact updates. Through the lens of reproducing kernel
Hilbert spaces (RKHS), UniCon provides a kernelized perspective that unifies contrastive
alignment and reveals its connection to spectral methods. To validate the theory, we conduct experiments on synthetic, unimodal,
multimodal, and zero-shot tasks, demonstrating that UniCon achieves substantial efficiency gains while preserving generality and strong empirical performance.
Read the full abstract ▾
Machine Learning
Cross-Attentive Bayesian Low-Rank Adaptation for Multimodal Uncertainty Estimation
Habibeh Naderi, Behrouz Haji Soleimani and Stan Matwin
Parameter-efficient fine-tuning (PEFT) enables efficient adaptation of frozen
language models, but most PEFT methods remain deterministic or unimodal, limiting their reliability in low-resource audio-text settings where uncertainty depends on both linguistic evidence and acoustic conditions. We introduce SPECTRA (
Stochastic Posterior Estimation with Cross-modal
Token-level Rank Adaptation), a
multimodal Bayesian low-rank adaptation framework for uncertainty-aware audio-text learning. SPECTRA keeps the text and audio backbones frozen and confines stochasticity to a compact rank-$r$ latent
matrix inside each LoRA adapter. At each
transformer layer, text-derived low-rank
token features query frame-level audio
embeddings through lightweight cross-
attention; the resulting
token-specific acoustic context parameterizes the mean and variance of an amortized variational posterior over the adapter latent. Audio thus acts not merely as an additional feature stream, but as a localized reliability signal that modulates both adaptation and confidence. Posterior prediction uses Monte Carlo adapter samples, enabling a decomposition of predictive uncertainty into total, aleatoric, and adapter-space epistemic components. Across IEMOCAP and clinical interview prediction tasks, SPECTRA is consistently competitive with or improves upon text-only
Bayesian PEFT and conventional
multimodal transfer-learning baselines, with
token-level cross-
attention yielding the most reliable gains. Modality-disagreement stress tests show that mismatched audio causes substantially larger degradation in AUC, likelihood, calibration, and Brier score than noisy but matched audio, highlighting the importance of modeling cross-modal reliability. These results suggest that
Bayesian cross-modal conditioning in low-rank adapter space provides an efficient and principled mechanism for calibrated
multimodal adaptation.
Read the full abstract ▾
Machine Learning
FLARE: A Data-Efficient Surrogate for Predicting Displacement Fields in Directed Energy Deposition
Kittipong Thiamchaiboonthawee, Ghadi Nehme, Ram Mohan Telikicherla et al. (+5 more)
Directed energy deposition (DED) produces complex thermo-mechanical responses that can lead to distortion and reduced dimensional accuracy of a manufactured part. Thermo-mechanical finite element simulations are widely used to estimate these effects, but their computational cost and the complexity of accurately capturing DED physics limit their use in design iteration and process optimization. This paper introduces
FLARE (Field Prediction via Linear Affine Reconstruction in wEight-space), a data-efficient surrogate modeling framework for predicting post-cooling displacement fields in DED from
geometric and process parameters. We develop a predefined-geometry DED simulation workflow using an open-source finite element framework and generate a
dataset of simulations with varying geometry,
laser power, and deposition velocity. Each simulation provides full-field displacement, stress, strain, and temperature data throughout the manufacturing process.
FLARE encodes each simulation as an implicit neural field and regularizes the corresponding neural-network weights so that they follow the affine structure of the input parameter space. This enables prediction of unseen parameter combinations by reconstructing network weights through affine mixing of training examples. On this DED
benchmark, the method shows improved accuracy compared to baseline methods in both in-distribution and extrapolation settings. Although the present study focuses on DED displacement prediction, the proposed affine weight-space reconstruction framework offers a promising approach for data-efficient surrogate modeling of physical fields.
Read the full abstract ▾
All authors Kittipong Thiamchaiboonthawee, Ghadi Nehme, Ram Mohan Telikicherla, Jiawei Tian, Balaji Jayaraman, Vikas Chandan, Dhanushkodi Mariappan, Faez Ahmed
Machine Learning
FRIGID: Scaling Diffusion-Based Molecular Generation from Mass Spectra at Training and Inference Time
Montgomery Bohde, Hongxuan Liu, Mrunali Manjrekar et al. (+4 more)
Tandem mass spectrometry is prominent in scientific discovery workflows for identifying unknown small molecules, yet high-throughput structural elucidation remains challenging. While recent autoregressive and graph
diffusion models have shown promise in de novo elucidation, performance remains limited by poor scalability during both training and inference time. In this work, we present FRIGID, a framework with a novel
diffusion language model that generates molecular structures conditioned on mass spectra via intermediate fingerprint representations and determined chemical formulae, training at the scale of hundreds of millions of unlabeled structures. We then demonstrate how forward fragmentation models enable inference-time scaling by identifying spectrum-inconsistent fragments and refining them through targeted remasking and denoising. While FRIGID already achieves strong performance with its
diffusion base, inference-time scaling significantly improves its accuracy, surpassing 18% Top-1 accuracy on the challenging MassSpecGym
benchmark and tripling the Top-1 accuracy of the leading methods on NPLIB1. Further empirical analyses show that FRIGID exhibits log-linear performance scaling with increasing inference-time compute, opening a promising new direction for continued improvements in de novo structural elucidation. FRIGID code is publicly available at https://github.com/coleygroup/FRIGID.
Read the full abstract ▾
All authors Montgomery Bohde, Hongxuan Liu, Mrunali Manjrekar, Magdalena Lederbauer, Shuiwang Ji, Runzhong Wang, Connor W. Coley
Machine Learning · Cross-listed from cs.CL
AdaExplore: Failure-Driven Adaptation and Diversity-Preserving Search for Efficient Kernel Generation
Weihua Du, Jingming Zhuo, Yixin Dong et al. (+9 more)
Recent large
language model (LLM)
agents have shown promise in using execution feedback for test-time adaptation. However, robust self-improvement remains far from solved: most approaches still treat each problem instance independently, without accumulating reusable knowledge. This limitation is particularly pronounced in domain-specific
languages such as Triton, which are underrepresented in LLM pretraining data. Their strict constraints and non-linear optimization landscape further make naive generation and local refinement unreliable. We propose AdaExplore, an
agent framework that enables self-improvement via accumulated execution feedback for performance-critical kernel code generation through two complementary stages: failure-driven adaptation and diversity-preserving search, jointly improving correctness and optimization performance without additional fine-tuning or external knowledge. In the adaptation stage, the
agent synthesizes tasks and converts recurring failures into a reusable
memory of validity rules, helping subsequent generations remain within the feasible set. In the search stage, the
agent organizes candidate kernels as a tree and alternates between small local refinements and larger structural regeneration, allowing it to explore the optimization landscape beyond local optima. Experiments on kernel runtime optimization
benchmarks validate these gains: AdaExplore achieves 3.11x and 1.62x speedups on KernelBench Level-2 and Level-3, respectively, within 100 steps, and continues to improve with additional computation.
Read the full abstract ▾
All authors Weihua Du, Jingming Zhuo, Yixin Dong, Andre Wang He, Weiwei Sun, Zeyu Zheng, Manupa Karunaratne, Ivan Fox, Tim Dettmers, Tianqi Chen, Yiming Yang, Sean Welleck
Machine Learning · Cross-listed from cs.CL
Aligning Backchannel and Dialogue Context Representations via Contrastive LLM Fine-Tuning
Livia Qian and Gabriel Skantze
Backchannels (e.g., `yeah', `mhm', and `right') are short, non-interruptive feedback signals whose lexical form and prosody jointly convey pragmatic meaning. While prior computational research has largely focused on predicting backchannel timing, the relationship between lexico-prosodic form and meaning remains underexplored. We propose a two-stage framework: first, fine-tuning large
language models on dialogue transcripts to derive rich contextual representations; and second, learning a joint
embedding space for dialogue contexts and backchannel realizations. We evaluate
alignment with human perception via triadic similarity judgments (prosodic and cross-lexical) and a context-backchannel suitability task. Our results demonstrate that the learned projections substantially improve context-backchannel
retrieval compared to previous methods. In addition, they reveal that backchannel form is highly sensitive to extended conversational context and that the learned
embeddings align more closely with human judgments than raw WavLM features.
Read the full abstract ▾
Machine Learning
Lower Bounds and Proximally Anchored SGD for Non-Convex Minimization Under Unbounded Variance
Arda Fazla, Ege C. Kaya, Antesh Upadhyay et al. (+1 more)
Analysis of
Stochastic Gradient Descent (SGD) and its variants typically relies on the assumption of uniformly bounded variance, a condition that frequently fails in practical non-convex settings, such as neural network training, as well as in several elementary optimization settings. While several relaxations are explored in the literature, the Blum-Gladyshev (BG-0) condition, which permits the variance to grow quadratically with distance has recently been shown to be the weakest condition. However, the study of the oracle complexity of
stochastic first-order non-convex optimization under BG-0 has remained underexplored. In this paper, we address this gap and establish information-theoretic lower bounds, proving that finding an $\epsilon$-stationary point requires $\Omega(\epsilon^{-6})$
stochastic BG-0 oracle queries for smooth functions and $\Omega(\epsilon^{-4})$ queries under mean-square smoothness. These limits demonstrate an unavoidable degradation from classical bounded-variance complexities, i.e., $\Omega(\epsilon^{-4})$ and $\Omega(\epsilon^{-3})$ for smooth and mean-square smooth cases, respectively. To match these lower bounds, we consider Proximally Anchored
STochastic Approximation (PASTA), a unified algorithmic framework that couples Halpern anchoring with Tikhonov regularization to dynamically mitigate the extra variance explosion term permitted by the BG-0 oracle. We prove that PASTA achieves minimax optimal complexities across numerous non-convex regimes, including standard smooth, mean-square smooth, weakly convex,
star-convex, and Polyak-Lojasiewicz functions, entirely under an unbounded domain and unbounded
stochastic gradients.
Read the full abstract ▾
All authors Arda Fazla, Ege C. Kaya, Antesh Upadhyay, Abolfazl Hashemi