Machine Learning
AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis
Jaber Jaber and Osama Jaber
AutoMegaKernel (AMK) compiles a HuggingFace Llama-family model into a single persistent cooperative CUDA kernel that runs the whole forward pass in one launch, with no per-model hand-written CUDA. The contribution is the system, not raw speed. A frozen schedule-IR validator statically certifies deadlock-freedom and race-freedom via static graph checks (not a mechanized proof), so an unsafe
agent-proposed schedule is rejected before launch: across 7,160
adversarial schedules (6,091 unsafe) it had zero false-accepts and accepted all 360 real lowerings. The same source retargets sm_80/sm_90/sm_120 from one codebase, auto-generates correct megakernels for 10 of 10 supported models, and on a real SmolLM2-135M checkpoint reproduces HuggingFace greedy decode
token-for-
token (perplexity match 2.5e-7). An unattended,
agent-drivable autoresearch loop self-improves the megakernel over its own baseline (1.25-1.72x). A search-found int8 (W8A16) megakernel beats CUDA-graphed cuBLAS bf16 at batch-1 decode across NVIDIA's datacenter inference fleet: L4 up to 1.33x, the current-gen L40S 1.25-1.27x, A10G up to 1.08x at scale, and the consumer RTX 5090 1.19-1.23x. The ordering is not a clean function of bandwidth (the 864 GB/s L40S beats the 600 GB/s A10G); the divide is inference-class vs training-class. AMK trails cuBLAS on the high-bandwidth training-class A100/H100, where the harness localizes the cross-SM-sync bottleneck; we report the gap plainly. This is a precision-asymmetric (W8A16 vs bf16) comparison at decode position 0; the largest real checkpoint is TinyLlama-1.1B. Code and the harness: https://github.com/RightNow-AI/AutoMegaKernel
Read the full abstract ▾
Machine Learning · Cross-listed from cs.AI
Correlation Is Not Enough: Embedding Human Metadata for Individual Causal Discovery
Suraj Biswas, Saurabh Gupta and Pritam Mukherjee
Ask a pretrained biomedical
language model whether "cortisol 28 ug/dL" and "stock-
market volatility" are related, and it returns a cosine similarity of 0.83 on a scale where 1.0 means identical. The two share no mechanism. This is not a corner case: every off-the-shelf biomedical encoder we tested (BioBERT, PubMedBERT, BioM-ELECTRA) scores unrelated cross-domain pairs between 0.76 and 0.92 when the answer should be near zero. Accuracy on cross-domain discrimination is 0%.
Retrieval systems survive this, because a
language model downstream filters the noise. A Large Behavioural Model (LBM), a foundation model whose subject is a person rather than a sentence, does not: it reasons over a graph of a user's life and treats
embedding proximity as evidence that two events are causally linked. False proximity writes a false causal edge, and everything downstream inherits the error. Here,
embedding geometry is not a tuning knob; it is correctness. We report the fix. A contrastive pass over 72,034 pairs raises PubMedBERT BIOSSES correlation from 0.633 to 0.828 and within-vs-across-domain separation from 1.05x to 1.63x. A second pass, BODHI, mines hard negatives from edges absent in a biomedical knowledge graph and lifts separation to 2.30x and the discrimination gap to +0.392, at a 4.5% BIOSSES cost. On an Intel Xeon 6737P with AMX, OpenVINO cuts single-query latency from 1367 ms to 10 ms (133x) and reaches 555 sentences/sec. One finding contradicts standard advice: FP16 beats INT8 on this silicon at every serving batch size, and we explain why. The same model on a no-AMX Ice Lake instance runs 13-27x slower. We release the
benchmark suite, training corpora, the BODHI generator, and the OpenVINO scripts.
Read the full abstract ▾
Machine Learning
Transition-Based Digital Twin Modelling for Alzheimer's Disease under Sparse Longitudinal Data
Yinyu Huang, Yilin Zhang, Sofia Michopoulou et al. (+2 more)
Alzheimer's disease (AD) progression is highly heterogeneous and is typically observed through sparse and irregular longitudinal data, posing challenges for prediction and personalised monitoring. Existing machine learning approaches have improved AD prediction using
multimodal data, yet often focus on static classification or cohort-level risk estimation, providing limited support for subject-specific modelling and uncertainty-aware
reasoning. To address these limitations, we present a personalised digital twin framework for AD prediction and scenario-based analysis using
multimodal longitudinal data. The proposed approach integrates complementary modelling strategies to capture clinical transitions and temporal dependencies across visits. Using data from the Alzheimer's Disease Neuroimaging Initiative (ADNI), including cognitive assessments, clinical variables, and MRI-derived phenotypes, the framework predicts cognitive status and diagnostic categories while quantifying predictive uncertainty and enabling patient-specific what-if trajectory analysis. Evaluation on leak-free subject-level splits demonstrates strong performance in score
forecasting and diagnosis classification. In this sparse and irregular ADNI setting, transition-based modelling of adjacent visits achieved higher predictive accuracy than the sequence-based branch, suggesting that local transition modelling may be more data-efficient. While sequence models remain valuable for uncertainty-aware trajectory
forecasting, local transition modelling offers a more data-efficient and robust predictive strategy. These findings highlight the importance of aligning temporal modelling strategies with clinical data structure and suggest that transition-based digital twin formulations may provide a practical and interpretable approach for personalised disease
forecasting in neurodegenerative disorders.
Read the full abstract ▾
All authors Yinyu Huang, Yilin Zhang, Sofia Michopoulou, Christopher Kipps, Rahman Attar
Machine Learning
Algorithm for Contextual Queueing Bandits with Rate-Optimal Queue Length Regret
Seoungbin Bae and Dabeen Lee
Contextual queueing bandits provide a framework for learning to schedule heterogeneous jobs under unknown context-dependent service rates. Under
stochastic contexts, existing
algorithms achieve $\widetilde{\mathcal{O}}(T^{-1/4})$ queue length regret, defined as the expected difference between the learner's and oracle's queue lengths at horizon $T$. In this paper, we improve this rate to $\widetilde{\mathcal{O}}(T^{-1/2})$. The key observation is that random exploration is needed only up to a carefully chosen cutoff round, rather than throughout the entire horizon. We propose CQB-$\eta$-2, a three-phase
algorithm: (i) pure random exploration to construct an initial
estimator, (ii) $\eta$-random exploration combined with a UCB rule to continue learning while maintaining negative drift, and (iii) pure UCB after the exploration cutoff. Our proof decomposes the queue length regret at the cutoff round. Before the cutoff, negative drift suppresses queue length differences caused by suboptimal choices. After the cutoff, the first two phases provide sufficient random exploration samples, ensuring that UCB decisions incur small departure-rate gaps. Combining these two bounds yields queue length regret of order $\widetilde{\mathcal{O}}(T^{-1/2})$. We further prove a minimax lower bound of order $\Omega(T^{-1/2})$. The proof constructs two hard instances that are statistically indistinguishable up to the final service decision, and uses a queue-specific coupling argument to convert the resulting testing error into queue length regret. Together, our upper and lower bounds characterize the minimax dependence on the horizon $T$ up to logarithmic factors.
Read the full abstract ▾
Machine Learning
In-Context Learning for Latent Space Bayesian Optimization
Tuan A. Vu, Harri Lähdesmäki and Julien Martinelli
Bayesian optimization (BO) is a central tool for sample-efficient design, and latent-space
Bayesian optimization (LSBO) extends it to structured objects such as molecules and
proteins. In parallel, tabular foundation models such as TabPFN and TabICL now achieve state-of-the-art regression performance and are increasingly used as BO surrogates. Because their
Bayesian behavior is induced by large synthetic pretraining collections, the composition of this pretraining distribution is crucial. LSBO creates a distinctive mismatch: the induced map from latent code to objective value differs markedly from the regression tasks used to train current in-context models. We address this mismatch by complementing the pretraining stage of tabular foundation model surrogates with synthetic optimization tasks defined on the latent space of a molecular VAE. The continued-pretraining objective features a regularizer that anchors the model to the original checkpoint, preserving its broad regression prior while avoiding overspecialization to the adaptation tasks. On held-out molecular optimization
benchmarks, the resulting model achieves strong performance, supporting the relevance of LSBO-specific adaptation for in-context surrogates.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
End-to-End Context Compression at Scale
Ang Li, Sean McLeish, Haozhe Chen et al. (+12 more)
Long-context
language model inference is bottlenecked by
memory, as the KV cache grows with context length. Recent techniques to compress the KV cache fall short: they either degrade model quality substantially or require considerable time and compute to compress a single long
prompt. Furthermore, many methods require the input to fit within the target model's context window, and are generally incompatible with modern production inference engines. Encoder-decoder compressors, which map a long
token sequence to a shorter sequence of latent
embeddings consumed by a decoder, are an appealing alternative in principle. However, existing approaches are not competitive with KV cache compression on the accuracy-efficiency frontier. In this work, we revisit encoder-decoder compression and close this gap. We first perform an architecture search, pre-training many variants from scratch to determine how best to design and train encoder-decoder compressors. Guided by our findings, we continually pre-train a family of 0.6B-encoder, 4B-decoder models on nearly 400B
tokens each, at compression ratios of 1:4, 1:8, and 1:16. We introduce Latent Context
Language Models (LCLMs), a family of compressors that improve the Pareto frontier across general-task performance, compression speed, and peak
memory usage. LCLMs also unlock new efficient downstream applications: they serve as backbones for long-horizon
agents that skim a compressed long context and adaptively expand relevant segments on demand, and they can compact the model's own generated
reasoning into latent
memory during inference, which is nearly lossless on the full MATH test set while being highly efficient, reducing peak KV usage by up to $7.9\times$.
Read the full abstract ▾
All authors Ang Li, Sean McLeish, Haozhe Chen, Nimit Kalra, Zaiqian Chen, Artem Gazizov, Venkata Anoop Suhas Kumar Morisetty, Bhavya Kailkhura, Harshitha Menon, Zhuang Liu, Brian R. Bartoldson, Tom Goldstein, Sanae Lotfi, Micah Goldblum, Pavel Izmailov
Machine Learning
Muon Learns More Robust and Transferable Features than Adam
Tianyu Ruan, Fengzhuo Zhang, Shuche Wang et al. (+1 more)
Muon has recently emerged as a state-of-the-art optimizer for pretraining Large
Language Models (LLMs) and
vision classifiers. Despite its efficiency advantage over Adam and SGD, the feature-learning advantage of
Muon remains unclear. This paper investigates
Muon's feature-learning advantage through the lens of robustness and transferability. First, by evaluating pretrained models on corrupted images and texts, we show that features learned by
Muon are consistently more robust than those learned by Adam and SGD across different architectures, including
transformers and Convolutional Neural Networks (CNNs). Using trained layer-wise probes, we further show that this robustness advantage is reflected in larger logit margins across layers. Second, by training linear classifiers or fine-tuning full models from pretrained parameters on downstream tasks, we demonstrate that
Muon-learned features transfer more effectively than those learned by Adam and SGD. This transferability advantage is further supported by the diversity of hidden states across layers, as measured by effective rank. Finally, in a representative classification problem with multi-component features, we prove that
Muon attains larger margins and higher effective rank than Adam and SGD, providing theoretical support for our empirical findings.
Read the full abstract ▾
All authors Tianyu Ruan, Fengzhuo Zhang, Shuche Wang, Shihua Zhang
Machine Learning
A Unifying Framework for Concept-Based Representational Similarity
Grégoire Dhimoïla, Agustin Martin Picard, Victor Boutin et al. (+2 more)
Learned representations across models and modalities often exhibit striking structural similarities, suggesting shared underlying concept decompositions. However, concept
alignment remains poorly defined: existing approaches optimize different objectives under the same terminology, obscuring what is actually aligned. We propose a unifying framework that decomposes
alignment along two axes: what is aligned (representations vs. concepts) and at what level (instance-wise vs. distributional). This induces four corresponding properties—instance-wise and distributional variants of translation and concept consistency—and reveals precisely which of these guarantees existing methods provide. We further introduce InterVenchA, an intervention-based
benchmark that separately measures extraction quality, translation quality, and concept consistency—which we validate against human perception of concept
alignment. Using this framework, we show that commonly assumed equivalences between
alignment objectives fail in practice: optimizing one property does not reliably recover the others, and, on real
embeddings, purely unsupervised objectives fail to recover meaningful instance-level
alignment. As a constructive consequence, a model, the Coupled SAE (CoSAE), matches or exceeds specialized
alignment methods such as crosscoders and Universal SAEs on their own target metrics, by jointly enforcing the properties identified by our framework without any architectural innovation. Moreover, as little as 1% paired data is sufficient to recover instance-level
alignment when anchoring distributional objectives. Overall, concept
alignment is fundamentally multi-objective, and progress requires treating it as such. Our framework and
benchmark provide the means to define, measure, and compare concept
alignment methods on this basis.
Read the full abstract ▾
All authors Grégoire Dhimoïla, Agustin Martin Picard, Victor Boutin, Thomas Fel, Thomas Serre
Machine Learning · Cross-listed from cs.CV
How Do Video Foundation Models Encode Intuitive Physics? Probing Across Pretraining Paradigms
Samuele Punzo, Niccolò Caselli, Ippokratis Pantelidis et al. (+3 more)
Do
video foundation models encode a usable sense of how the physical world works, and does intuitive physics emerge naturally from large-scale
video pretraining? In this work, we study whether pretrained
video foundation models encode intuitive-physics information in their frozen representations, and how its accessibility varies across pretraining paradigms, layers, backbone variants, and probe types. Using frozen-feature probing on IntPhys2 and Minimal
Video Pairs (MVP), we compare predictive joint-
embedding models (V-JEPA), masked reconstruction models (VideoMAE), and
diffusion-based
video generators (LTX-
Video). We find that across linear, MLP, and temporal attentive probes, physics-relevant information is generally weakest in early layers and most accessible at intermediate to late depth. On MVP, temporal attentive probes substantially outperform linear readouts, with V-JEPA achieving 95.0% pair consistency. On IntPhys2, linear probes recover more signal, and VideoMAE attains the strongest temporal attentive result at 73.9% violation-of-expectation accuracy. On both
benchmarks, attentive probe performance on LTX consistently underperforms compared to the other two pretraining paradigms. Together, these results suggest that intuitive-physics is broadly decodable from
video foundation models, but it is not represented as a single, uniform capability. Its accessibility, and especially its temporal grounding, depends much more on the
benchmark and readout than on a simple hierarchy of pretraining objectives or model scale.
Read the full abstract ▾
All authors Samuele Punzo, Niccolò Caselli, Ippokratis Pantelidis, Francesco Massafra, Salvatore Lo Sardo, Mohammadreza Salehi
Machine Learning · Cross-listed from cs.DC
FMplex: Model Virtualization for Serving Extensible Foundation Models
Hetvi Shastri, Pragya Sharma, Walid A. Hanafy et al. (+3 more)
Foundation models (FMs) are increasingly used as backbones for downstream tasks across
language,
vision, time-series, and
multimodal applications. Yet existing model-serving systems deploy each customized task as an independent model instance, thereby replicating heavyweight backbones, wasting accelerator
memory, and losing opportunities to amortize batching and loading costs. This paper presents FMplex, a serving system that treats FM backbones as a virtualization substrate for deployment sharing. FMplex presents each task with a virtual foundation model (vFM), a logically private FM instance backed by a shared physical FM. This abstraction lets independently customized tasks share a backbone while preserving task-specific extensions, independent lifecycles, and task-level isolation. In addition, we propose a batch-aware fair-queueing scheduler that combines weighted task-level sharing with inter- and intra-task batching across colocated tasks. We implement a FMplex-based serving stack spanning task construction, sharing-aware deployment, and runtime execution. Across 7 FM backbones (16 variants) and 92 downstream tasks, FMplex reduces latency by up to 80% over spatial partitioning and 33.3% over best-effort co-location, while hosting up to 6x more tasks at cluster scale.
Read the full abstract ▾
All authors Hetvi Shastri, Pragya Sharma, Walid A. Hanafy, David Irwin, Mani Srivastava, Prashant Shenoy
Machine Learning
Data-driven discovery of governing differential equations across physical systems
Siyu Lou, Hao Xu, Wenguan Wang et al. (+6 more)
Differential equations play a critical role in scientific discovery because they provide a mathematical framework to describe the behaviour of physical phenomena. As a promising alternative to traditional first principles, data-driven differential equation discovery has attracted increasing
attention for its ability to infer governing laws directly from experimental or simulated data, especially when the underlying physics is unclear. However, the field has expanded rapidly along diverse methodological directions, particularly with the emergence of AI-based approaches, and still lacks a clear organizing perspective. In this Review, we propose a problem-oriented perspective on data-driven differential equation discovery. We first introduce a two-dimensional phase diagram of equation discoverability, where discovery problems are organized according to structural complexity and coefficient complexity. This phase diagram shows how the field has moved from the discovery of sparse equations with simple coefficients toward more complex governing laws with richer structures and more flexible parameterizations. It also clarifies why different methodological families succeed or fail in different problem settings. We then present the representation-evaluation-optimization (REO) framework as a fundamental abstraction of the discovery process. By identifying the core problems of equation discovery that persist across algorithmic variations, REO shifts the discussion from individual
algorithms to the fundamental principles that determine discoverability. We connect these perspectives to applications across physics and adjacent sciences, and argue that the next challenge is not merely recovering equations, but using them to revise existing theories, distil mechanisms and form new scientific concepts.
Read the full abstract ▾
All authors Siyu Lou, Hao Xu, Wenguan Wang, Lu Lu, Hao Sun, Yang Liu, Linfeng Zhang, Dongxiao Zhang, Yuntian Chen
Machine Learning · Cross-listed from cs.CL
Gradient-Guided Reward Optimization for Inference-time Alignment
Hankun Lin and Ruqi Zhang
Ensuring the reliability of Large
Language Models (LLMs) under distribution drift requires inference-time adaptation. While inference-time
alignment methods such as Best-of-$N$ and rejection sampling are widely used, they frame the task as a sampling-intensive, reward-guided search, leading to two key limitations: their performance is bounded by the base model's generation quality, and their reliance on imperfect reward models makes them vulnerable to reward hacking. To address these challenges, we introduce Gradient-Guided Reward Optimization (GGRO), a lightweight inference-time method that performs targeted, minimal intervention during decoding via gradient guidance. Specifically, GGRO monitors
token-level
entropy to identify high-uncertainty regions indicative of drift or misalignment. Upon detection, it responds by injecting nudging
tokens, generated using gradient signals from an off-the-shelf reward model, to steer the generation trajectory rather than merely re-ranking samples. Experiments show that GGRO consistently improves inference-time
alignment across safety, helpfulness, and
reasoning benchmarks. It also increases coverage of high-quality responses and robustness to reward hacking, with minimal computational overhead. Code is available at https://github.com/lhk2004/GGRO.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.RO
ReCoVLA: VLM-Guided Reward Compilation for Failure Recovery in Vision-Language-Action Policies
Haodi Hu, Chung-Ta Huang, Jing Liu et al. (+4 more)
Vision-
language-action (VLA) policies provide strong priors for
language-conditioned manipulation, but remain brittle in off-nominal states requiring targeted recovery. We propose ReCoVLA – a failure-conditioned residual recovery framework that keeps a pretrained VLA policy frozen, uses an external
vision-
language model (VLM) to infer the failure mode and recovery stage, and compiles a structured reward from task-relevant components. Rather than using the VLM to generate actions or rewards directly, ReCoVLA uses it as a semantic reward selector: it predicts a recovery descriptor and reward mask for in-simulation residual-policy training, followed by zero-shot sim-to-real deployment of the trained recovery policies. This decouples high-level failure understanding from low-level corrective control to support different VLAs. Experiments across short-horizon, long-horizon, and contact-rich manipulation tasks show that ReCoVLA outperforms the tested baselines on average. In simulation, our reward compiler improves average success from 36.7% for the fine-tuned $\pi_{0.5}$ baseline to 66.7%. In physical zero-shot sim-to-real experiments, ReCoVLA achieves the best average performance, with 61.7% success.
Read the full abstract ▾
All authors Haodi Hu, Chung-Ta Huang, Jing Liu, Ye Wang, Kei Suzuki, Matthew Brand, Toshiaki Koike-Akino
Machine Learning
Joint user-item allocation for e-commerce marketing campaigns
Maja Lindström, Natalija Glisovic, Jan von Pichowski et al. (+2 more)
Web platforms run many marketing campaigns at once, and for each one they must decide which products to promote and which users to target. The best users depend on the products, and the best products on the users. The campaigns also compete with one another, because each user and each item may belong to only one of them and each audience is capped by the campaign budget. Still, existing approaches fix one side and search the other, choosing an assortment and then looking for its audience, or ranking items for one user at a time. Either way they solve a restricted version of the full problem, implying that a jointly optimized assignment would be at least as good and possibly far better. What the restriction costs, and whether the joint problem is solvable at the scale campaigns are run, remain unknown. Here we formalize the joint task as auto-targeting: assigning users and items to multiple disjoint, capacity-bounded campaigns that maximize the affinity between them. We adapt constrained spectral biclustering, greedy local search with pairwise swaps, and a multi-armed bandit formulation to the task and compare them against the conventional decoupled baseline. On synthetic data, two Amazon Reviews
benchmarks, and two proprietary commercial
datasets, auto-targeting raises campaign lift over the best decoupled pipeline, from around 20% on the two smaller, denser
datasets to more than threefold on the two largest. Biclustering yields the highest campaign quality across
datasets, while bandit-based methods run fastest on the largest ones.
Read the full abstract ▾
All authors Maja Lindström, Natalija Glisovic, Jan von Pichowski, Tommy Löfstedt, Martin Rosvall
Machine Learning
Closure-Validated Circuit Discovery in Attention Heads: Co-activation Proposes, Ablation Disposes
Yongzhong Xu
Interpretability increasingly treats groups of components, not individual units, as the basic object, and proposes to find them by clustering co-activation statistics. We ask whether such a cheap signal actually identifies an
attention-head
circuit. Adapting a sparse-autoencoder clustering recipe to
attention heads – but validating by causal ablation rather than reconstruction – we cluster heads and then run a closure test: ablate the discovered community and compare per-example damage to matched-random controls. Across two dense 1B-scale models (Pythia 1B, OLMo 1B) and two input distributions, the communities pass closure. In a Mixture-of-Experts model (OLMoE-1B-7B), route-conditional clustering recovers a statistically real signal that nonetheless does not survive closure – ablation improves loss, the wrong direction. Extending closure across training,
attention-target selectivity and participation ratio decouple from function in both directions. We conclude that a cheap signal is a
circuit proposal, not a confirmed
circuit; closure is what separates them.
Read the full abstract ▾
Machine Learning
Assessing Sample Quality in Conditional Generation under Compositional Shift
Berker Demirel, Valentino Maiorca, Marco Fumero et al. (+2 more)
Conditional generators provide a natural tool for controllable generation, including settings where the desired condition is a new composition of observed attributes or experimental factors. In many applications, especially in scientific domains, such models are attractive to explore conditions for which real samples are rare, expensive, or not yet observed. However, this creates a circularity for evaluation: standard conditional quality metrics require a reference target distribution, but in the extrapolative regime that distribution is unavailable by definition. We address this problem with a post-hoc, per-sample trust score for assessing conditional samples using only the training distribution. The score combines two estimable quantities: global realism, measuring compatibility with the real data
manifold, and attribute-wise faithfulness, measuring whether a sample is closer to the requested attributes than to plausible alternatives. We show that the score can recover meaningful comparisons across extrapolated generations, under a mild coverage condition on the observed attributes. These comparisons enable effective filtering, ranking, and abstention of generations and can be used directly on off-the-shelf pretrained models. In biological imaging, selected samples preserve real morphological structure better and improve downstream predictive performance, while similar gains are observed on controlled
vision benchmarks. Finally, we show how the score can be applied during generation, enabling abstention before full decoding. Code is available at https://github.com/berkerdemirel/faithful-cond-gen.
Read the full abstract ▾
All authors Berker Demirel, Valentino Maiorca, Marco Fumero, Theofanis Karaletsos, Francesco Locatello
Machine Learning
On Choosing the $\mu$ Parameter in Gaussian Differential Privacy
Bogdan Kulynych and Antti Honkela
Recent work argues for using Gaussian differential
privacy (GDP) to report the
privacy guarantees in
privacy-preserving machine learning. We provide principled mappings from pure-DP $\varepsilon$ to GDP $\mu$ by matching the worst-case success of a strong-adversary membership inference attack in terms of three metrics: multiplicative advantage at fixed FPR, precision at fixed recall, and the standard
privacy profile. We tabulate $\mu$ values across a useful range of parameters and recommend $\mu \approx \varepsilon/5$ as a conservative general-purpose conversion.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
Code Is More Than Text: Uncertainty Estimation for Code Generation
Yuling Shi, Caiqi Zhang, Yuexian Li et al. (+4 more)
Large
language models (LLMs) are increasingly deployed as code generators, where silently wrong programs pose real safety and reliability risks. Reliable uncertainty estimation (UE) is essential for selective prediction, human-in-the-loop review, and downstream agentic decisions. Yet most existing code UE methods are inherited from natural
language (NL) generation and ignore properties that make code distinct. We argue that code differs from NL in three ways: a single wrong
token can break an entire program (
token fragility); algorithmic intent and concrete implementation can disagree independently (intent-code gap); and programs can be executed (executability). We instantiate these properties as three orthogonal uncertainty axes: lexical (Top-K
token entropy), algorithmic (pseudo-code consistency), and functional (behavioral consistency). Across five code LLMs, our three-axis ensemble improves average AUROC from 0.696 for the strongest NL-derived baseline to 0.776 (+8.1 points). Notably, on Qwen3-14B, our single-pass Top-K
token entropy matches the strongest multi-pass baseline while being over 3x cheaper; across models, it remains a competitive low-cost signal. These results suggest that code UE deserves code-specific design rather than direct NL ports.
Read the full abstract ▾
All authors Yuling Shi, Caiqi Zhang, Yuexian Li, Haopeng Wang, Yeheng Chen, Nigel Collier, Xiaodong Gu
Machine Learning · Cross-listed from cs.AI
PRISM: Recovering Instruction Sets from Language Model Activations
Gilad Gressel, Rahul Pankajakshan, Julia Diament et al. (+3 more)
As LLMs are deployed as
agents, reliable monitoring requires knowing not only what they output, but which instructions are steering their behavior. This is difficult when models infer unintended subgoals, follow contextual cues, or are influenced by
prompt injections and hidden objectives. While activation-to-
language methods suggest that hidden states can reveal natural-
language information, existing approaches are not designed to recover the full set of simultaneous instructions, constraints, prohibitions, and subgoals present in agentic settings. As a first step toward this monitoring goal, we formalize instruction set
retrieval and introduce PRISM, an activation-conditioned interpreter that decodes hidden states from a frozen target model into a faithful bullet list of instructions. Unlike prior activation-to-
language methods, PRISM is trained to recover instruction sets directly, using judge-guided GRPO to reward covered instructions and penalize unsupported ones. Across benign, constrained,
prompt-injection, and hidden-objective settings, PRISM outperforms activation-to-
language baselines, especially on security-relevant objectives.
Read the full abstract ▾
All authors Gilad Gressel, Rahul Pankajakshan, Julia Diament, Efim Hudis, Krishnashree Achuthan, Yisroel Mirsky
Machine Learning
Safe-RULE: Safe Reinforcement UnLEarning
Shixiong Jiang, Taozheng Zhu and Fanxin Kong
Offline safe reinforcement learning (Safe RL) enables policy learning without online interactions, making it suitable for safety-critical systems such as robotics systems. However, its reliance on static
datasets exposes offline Safe RL to data poisoning attacks, where adversaries inject malicious samples that compromise safety and induce unsafe policy behavior. In this work, we propose a new learning paradigm, named safe reinforcement unlearning (Safe-RULE), used as a defense framework to remove the influence of poisoned data without retraining from scratch or requiring access to the original training environment. We further extend reinforcement unlearning to offline Safe RL by explicitly accounting for both task performance and safety constraints during the unlearning process. Experiments across
benchmark Safe RL tasks demonstrate that our approach effectively enhances safety performance against data poisoning attacks.
Read the full abstract ▾
Machine Learning · Cross-listed from q-bio.GN
Integrating gene regulatory priors into Transformer attention with scTransformer for interpretable scRNA-seq analysis
Mikele Milia, Louis Fabrice Tshimanga, Henning Mueller et al. (+2 more)
Motivation:
Transformer-based models are increasingly applied to large-scale single-
cell transcriptomics, showing strong performance through self-supervised learning on millions of
cells. However, most existing approaches treat genes as independent features, and largely ignore prior biological knowledge, which limits interpretability and robustness. In this paper, we explore whether explicitly incorporating gene regulatory information can improve both model performance and biological insight. Results: We present scTransformer, the first
Transformer-based approach that builds a priori knowledge of biological mechanisms into the model's
attention patterns. By constraining information flow according to known regulatory structures, the model learns representations that are more biologically meaningful. We evaluate scTransformer on a disease-relevant single-
nucleus RNA-seq
dataset using supervised
cell-type classification. Compared to standard
Transformers, our approach improves classification accuracy, enhances separation of
cell types in
embedding space, and produces
attention patterns consistent with known regulatory programs. Overall, our results demonstrate that
embedding biological structure into
Transformer models can enhance interpretability without sacrificing performance, offering a principled step toward biologically grounded foundation models for single-
cell omics.
Read the full abstract ▾
All authors Mikele Milia, Louis Fabrice Tshimanga, Henning Mueller, Manfredo Atzori, Barbara Di Camillo
Machine Learning · Cross-listed from cs.CV
Streaming Interventions: Can Video Large Language Models Correct Mistakes as They Occur?
Apratim Bhattacharyya, Shweta Mahajan, Sanjay Haresh et al. (+5 more)
Learning everyday skills, like cooking a dish, relies increasingly on instructional media such as online
videos. This opens the door to the use of
video (and
multimodal) large
language models (LLMs) as task guidance assistants. A crucial capability for the real-world success of a prospective task guidance assistant is it's ability to intervene proactively as soon as a mistake is apparent in order to guide the user. To evaluate this crucial capability, we introduce Ego-MC-Bench (Mistake Corrections), a
benchmark for evaluating reactive, step-by-step task guidance in realistic cooking scenarios. Extensive experiments show that Ego-MC-Bench is highly challenging for state-of-the-art
video LLMs. We argue that a key reason is the limited availability of training data for fine-tuning models on this task. Although there exists a wide range of cooking
video datasets, existing
datasets lack examples of mistakes along with appropriately timed interventions. To help address this data limitation, we also introduce Ego-CoMist, a counterfactual synthetic
dataset created by transforming non -interactive cooking
videos into supervised training examples showing proactive interventions. We show that fine-tuning on Ego-CoMist yields performance gains especially for smaller and more efficient
video LLMs that are well suited for delivering assistance on edge devices.
Read the full abstract ▾
All authors Apratim Bhattacharyya, Shweta Mahajan, Sanjay Haresh, Rajeev Yasarla, Reza Pourreza, Litian Liu, Risheek Garrepalli, Roland Memisevic
Machine Learning · Cross-listed from physics.app-ph
Automating the Expert Eye: A System-Agnostic Deep Learning Framework for Rare Event Discovery in Imbalanced Force Spectroscopy
Jorge Rodriguez-Ramos
Single-Molecule Force Spectroscopy (SMFS) provides unprecedented insights into biomolecular mechanics, yet the high-throughput generation of force-extension trajectories creates a severe data curation bottleneck. Identifying rare molecular unbinding events within thousands of noise-dominated curves traditionally relies on tedious, non-scalable manual
auditing. Here, we present a system-agnostic, interpretable deep learning framework tailored to overcome extreme class imbalance in automated SMFS triage. Utilizing 1D-to-2D rasterized
geometric matrices, we deployed a modified ResNet18 architecture governed by an asymmetric Focal Loss objective function. We evaluated this framework on the complex mechanical unfolding pathways of the R. champanellensis cellulosome. Under hyper-imbalanced test conditions where the target interaction constituted only 1.34% of the
dataset (13 true events out of 970 traces), the model achieved an overall accuracy of 0.9196 and a remarkable True Positive Rate (Recall) of 0.9231. By implementing an empirically calibrated dual-threshold triage system, the pipeline automatically discarded 880 unambiguous background noise traces , reducing the manual curation workload by over 90% while safely preserving high-value rare data. Finally, Gradient-weighted Class Activation Mapping (Grad-CAM) visually validated that the network's decisions are firmly anchored in the relevant
geometric features of the force curves, specifically localizing on the structural unbinding regions, effectively mitigating 'black-box' skepticism. Built for free cloud-based execution, this open-source tool democratizes scalable, highly precise molecular discovery across the biophysics community.
Read the full abstract ▾
Machine Learning
Efficient Traffic Prediction at Scale: A Systematic Study of STGCN Architectural Depth
Soban Nasir Lone, Mohamed Abouelela, Taeyoung Yu et al. (+2 more)
Spatio-temporal graph neural networks (STGNNs) have become the dominant approach for traffic prediction, yet their computational requirements pose challenges for practical deployment in intelligent transportation systems (ITS). While recent work has proposed efficient alternatives to STGNNs, a fundamental question remains unexplored: are these architectures themselves over-parameterised? We examine this question using the Spatio-Temporal Graph Convolutional Network (STGCN), one of the most widely adopted models in this domain. Through systematic experiments across four diverse traffic
datasets, we compare 1-block, 2-block (standard), and 3-block STGCN variants. Our findings reveal that the single-block architecture achieves optimal performance for short-term prediction (10 mins) on three of four
datasets, while incurring only marginal degradation ($\leq$1.8% relative error) at longer horizons. Crucially, the 2-block variant incurs 61% higher CPU inference latency and 37% lower throughput relative to 1-block – substantial overhead for resource-constrained ITS deployment. The 3-block architecture offers no favourable tradeoff, more than doubling computational cost for $<$0.5% relative improvement. These results suggest that the default 2-block STGCN may be over-parameterised for many applications, with implications for both practitioners deploying traffic prediction systems and researchers benchmarking efficiency-focused methods.
Read the full abstract ▾
All authors Soban Nasir Lone, Mohamed Abouelela, Taeyoung Yu, Jiwon Kim, Constantinos Antoniou
Machine Learning
Investigating Calibration Challenges in Probabilistic Electricity Price Forecasting
Jan Niklas Lettner, Hadeer El Ashhab and Benjamin Schäfer
As renewable energy integration increases
market volatility, probabilistic electricity price
forecasting has become essential for effective risk management. However, current-proper-scoring rules often prioritize forecast sharpness at the expense of calibration, leading to overconfident and statistically unreliable uncertainty estimates. This work highlights the critical gap between theoretical scoring and practical calibration, demonstrating that models can become mere proxies for deterministic forecasts when reliability is neglected. We conclude that future research must shift toward calibration-aware objectives and architectures to ensure the distributional integrity of energy
market forecasts.
Read the full abstract ▾
Machine Learning
BUDDY: BUdget-Driven DYnamic Depth Routing for Adaptive Large Language Model Inference
Yuhua Zhou, Shaoqi Yu, Shichao Weng et al. (+4 more)
Large
language models (LLMs) incur high inference cost due to their depth and parameter scale. Depth
pruning can reduce latency by skipping redundant
Transformer blocks, but existing methods (i) provide limited control under user-specific compute budgets and (ii) typically fix the routing path, failing to adapt as the context grows during decoding. We propose Buddy, a budget-driven dynamic depth routing framework. Buddy uses a lightweight Decision Module to score intermediate layers conditioned on the input and deterministically executes the top-k layers to satisfy a given budget. To support decode-time adaptation, Buddy reuses the first-layer KV cache as a low-overhead global context source and pools it together with the newest
token representation before each routing decision. When no explicit budget is provided, an optional Budget Predictor estimates an input-dependent compute level to balance quality and efficiency. Experiments on Llama-family and Qwen models show that Buddy is competitive with strong static
pruning baselines and often improves the accuracy-compute trade-off, while uniquely supporting strict budget control, decode-time rerouting, and multiple budgets within a single trained model.
Read the full abstract ▾
All authors Yuhua Zhou, Shaoqi Yu, Shichao Weng, Changhai Zhou, Mingze Yin, Fei Yang, Aimin Pan
Machine Learning
Loss-Guided Adaptive Scale Refinement for Molecular Force Prediction
Limin Yu
Molecular systems involve interactions across multiple spatial scales, from local coordination and short-range perturbations to long-range electrostatic and solvent-mediated effects. However, most molecular representation learning methods rely on manually predefined scales, and the task-optimal modeling scale may not coincide with these fixed levels. This study introduces a loss-guided adaptive scale refinement framework for molecular force prediction, treating predefined scales as initial anchors and discovering task-effective resolutions through interpolation, routing, differentiable scale updates, and scale pool refinement. Using a NaCl aqueous ionic system as a minimal testbed, this study constructs short-scale and long-range force prediction branches and analyzes their complementarity. Oracle hard routing reduces the overall force MAE from 399.65 to 382.67, while continuous oracle interpolation further reduces it to 380.96. In close-contact regimes with nearest-ion distance below 0.6 nm, the close-contact MAE decreases from 327.22 to 260.51. A minimal scale pool update experiment shows that starting from endpoint anchors {0,1}, loss-guided updates automatically generate intermediate scales and recover most of the continuous oracle performance. The final updated scale pool {0,0.125,0.25,0.375,0.5,0.75,1} achieves an overall MAE of 381.23. These results support adaptive scale refinement as a promising direction for molecular representation learning, especially when fixed-scale modeling is insufficient.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.AI
Emergent alignment and the projectability of ethical personas
Guillermo Del Pinal, Youngchan Lee, Calum McNamara et al. (+1 more)
Work on `emergent misalignment' shows that finetuning LLMs on narrow tasks can induce broadly misaligned behavior. This supports the `persona selection' (PSM) hypothesis: during pre-training, LLMs learn to simulate different characters/perspectives, which can be elicited and refined during post-training. This paper investigates the converse phenomenon, `emergent
alignment', and uses it to support and refine the PSM and motivate a novel desideratum for
alignment. We finetune a helpful-only model on broad and narrow safety tasks. To create SFT samples, we follow the `Constitutional AI' (CAI) approach and use four constitutions which encode reasonable
alignment strategies: deontology, consequentialism, virtue ethics, and aligning AIs as subordinate to human authority. For each, we show that finetuning on two narrow safety sub-categories reliably induces emergent
alignment over a representative set of general safety categories, and on safety subcategories that we directly filtered-out of the data sets used for narrow
alignment. To test the `PSM' using a more fine-grained evaluation, we used a multidimensional `ethical persona' diagnostic. For each constitutionally finetuned (broad/narrow) model, we evaluate how well their behavior matches their expected signature profile. Our results show that our CAI models acquire their expected “ethical persona” – e.g., the model narrowly fine-tuned on SFT samples created using the consequentialist constitution agrees significantly more with utilitarian than deontological beliefs. Yet our coarse and fine-grained evaluations show that there are significant differences across our (broad/narrow) finetuned CAI models in how well they project. We conclude that
alignment strategies should be evaluated, not just on their (in-distribution) general safety performance, but also specifically on their degree of coarse and fine-grained projectability.
Read the full abstract ▾
All authors Guillermo Del Pinal, Youngchan Lee, Calum McNamara, Alejandro Perez Carballo
Machine Learning · Cross-listed from stat.ML
Report the Floor: A Training-Free Conformal Interval Is a Mandatory Baseline for Probabilistic Time-Series Forecasting
Valery Manokhin
Probabilistic forecasters are increasingly learned, yet the baselines they are compared against are often weak or omitted. We show that the simplest possible conformal interval - a last-value point forecast wrapped in a finite-sample split-conformal residual
quantile, with no parameters and no training - is a far stronger baseline than its near-total absence from recent learned-
forecasting and conformal-time-series comparisons would suggest. In one-step-ahead online
forecasting across 2,217 real series from nine public sources (Monash, LOTSA, the LTSF traffic/electricity/weather suites, METR-LA, BOOM, nips/probts), this ConformalNaive interval decisively beats the naive value-
quantile baselines, the entire NPTS family (NPTS 73%, SeasonalNPTS 64% of series), and the published Conformal Seasonal Pools (CSP) method (71% of series, bootstrap 95% CI [69,73], paired Wilcoxon p approx 7.6e-135); it is on par with the simpler learned conformal predictors (RCI,
quantile regression; median relative Winkler within 2%) and is beaten only by the adaptive-online and ensemble methods (SPCI, ACI, AgACI), which track distribution shift and lead by 9-33% relative Winkler. It is also better calibrated than a trained neural forecaster: on the six
datasets that introduced DeepNPTS, the trivial floors cover the truth 84-85% of the time at a nominal 95%, versus DeepNPTS's 66%. At multi-step seasonal horizons the picture inverts: the random-walk floor is the weakest method and the seasonal pool (CSP) wins - a boundary we map. Finally we give ConformalNaive+, a one-line, training-free, horizon-adaptive selector that attains the better of two complementary floors at every horizon with restored coverage. We argue the matching conformal naive floor must be a mandatory baseline whenever a learned probabilistic forecaster claims gains.
Read the full abstract ▾
Machine Learning
Escaping the KL Agreement Trap in On-Policy Distillation
Haoran Xin, Anhao Zhao, Ying Sun et al. (+3 more)
On-policy
distillation (OPD) provides dense
token-level supervision by asking a teacher to score student-generated rollouts. However, when the student drifts into an unrecoverable prefix, the teacher may locally agree with the degraded state, producing low reverse KL but little corrective training signal. We identify this persistent regime as a low-KL agreement trap. Further analyses show that
tokens during and after such traps produce less useful supervision signals. We propose KAT (KL Agreement Trap Termination), an online OPD termination rule that detects persistent low-KL agreement with a dynamic training-adaptive threshold. By filtering weak supervision from degenerate agreement, KAT improves avg@k accuracy by 2.66% and pass@k by 3.43% across four mathematical
benchmarks, while reducing average rollout length by 59.73%.
Read the full abstract ▾
All authors Haoran Xin, Anhao Zhao, Ying Sun, Jin Li, Xiaoyu Shen, Hui Xiong