Machine Learning
Evaluation of Active Feature Acquisition Policies with Tabular Foundation Models
Yuta Kobayashi, Divyam Madaan and Shalmali Joshi
Active feature acquisition learns policies that sequentially acquire features to maximize information about a target variable. We study how to learn and evaluate such policies from finite offline data using prior-data fitted networks (PFNs), which are off-the-shelf models that output posterior predictive distributions without task-specific training. We show that under the imbalanced coverage of offline data, using total predictive
entropy as a reward creates an epistemic bias that penalizes acquiring sparsely observed features. Specifically, this reward conflates epistemic uncertainty (arising from lack of offline data) with aleatoric uncertainty (arising from uninformative features). To address this, we target the posterior expected (aleatoric)
entropy instead of the total predictive
entropy output by a PFN for evaluating feature acquisitions. Empirical evaluations on synthetic and real-world
datasets demonstrate that our approach consistently reduces value estimation bias and yields credible intervals with strong empirical coverage, which can translate to improved downstream policy selection.
Read the full abstract ▾
Machine Learning
What pass@k Cannot Measure: Evaluating Diversity and Capability Retention after Post-Training
Subham Rath, Raj Dandekar, Rajat Dandekar et al. (+1 more)
pass@$k$, the fraction of problems a model solves within $k$ sampled attempts, is the field's default protocol for deciding whether reinforcement-learning (RL) post-training on verifiable rewards improved a model. At the population level, pass@$k$ depends only on a problem's probability of a correct sample, with no term for how it is distributed across outputs. We show this gap is not academic. Training Qwen2.5-1.5B-Instruct on grade-school math with Group Relative Policy Optimization (GRPO) and with rejection-sampling fine-tuning (RFT, training on the model's own shortest verifier-passed rollout) moves three complementary diversity measures (
token-level
entropy, answer-level
entropy, unique answers per
prompt) in opposite directions, with zero overlap across three seeds per arm. The gap survives restricting to verifier-correct completions only (lexical diversity among correct solutions is 15% lower for GRPO, after controlling for length) and a count-controlled check isolating diversity among incorrect answers alone, ruling out that GRPO's higher accuracy alone explains it. Yet pass@8 and pass@32 show no consistent winner on GSM8K, and a hard MATH-500 subset shows the same pattern: separation only at low $k$. Compared against the starting checkpoint, no trained arm significantly improves hard-problem coverage: RFT is significantly worse, while GRPO is statistically indistinguishable from it - so GRPO's pass@1 edge over RFT reflects a smaller loss relative to Base, not a capability gain, a missing-control issue, not a failure of pass@$k$. On GSM8K, only pass@1, with no role in detecting diversity by construction, separates the arms cleanly, rewarding the arm whose correct solutions are least diverse. We argue this is a concrete instance of a standard evaluation protocol missing a property it is routinely used to certify.
Read the full abstract ▾
All authors Subham Rath, Raj Dandekar, Rajat Dandekar, Sreedath Panat
Machine Learning
Fed-BRDECS: Privacy-Preserving and Heterogeneity-Aware Federated Deep Embedded Clustering
Haemin Park, Diego Klabjan, Martin W. Braun et al. (+2 more)
Federated deep clustering seeks to learn clustering-friendly representations from decentralized unlabeled data while preserving client
privacy. However, Deep Embedded Clustering (DEC)-style objectives depend on global soft-assignment statistics that require clients to reveal their sensitive information. We propose Fed-BRDECS, a
privacy-preserving and heterogeneity-aware
federated deep embedded clustering framework. Fed-BRDECS replaces the globally normalized clustering objective with a locally computable sample-stability loss, avoiding the transmission of local soft-assignment distributions. To tackle non-IID client distributions, we introduce prediction-balanced sampling, which oversamples locally rare predicted clusters without requiring ground-truth labels, and centroid-level restarting, which periodically refreshes biased or inactive centroids. Experiments on image and text clustering
benchmarks show that Fed-BRDECS consistently outperforms representative
federated clustering and deep clustering baselines under both IID and non-IID partitions. We further demonstrate its applicability to
federated time-series anomaly detection, where it improves reconstruction-based detectors without adding inference-time cost.
Read the full abstract ▾
All authors Haemin Park, Diego Klabjan, Martin W. Braun, Xiuqi Li, Balakrishnan Ananthanarayanan
Machine Learning · Cross-listed from stat.ME
A perspective note on likelihood approximation and inference for complex simulation models using a chain of aggregated normalizing flows
Getachew K Befekadu
We present a new perspective on the problem of likelihood approximation within the framework of simulation-based inference that promotes scalable and controllable simulation routines for large-scale data analysis, allows efficient parameter space exploration or smooth interpolation in high-dimensions and, thus, supports valid statistical treatments of hypothesis testings as well as uncertainty quantification. In particular, we consider a chain of $n$-aggregated normalizing flows for likelihood approximation scheme, where a set of upfront replicated observation
datasets from the forward complex simulation model pass through the first set of bijective transformations, and then subsequently pass to the other sets of bijective transformations. Here, we assume that, for any $k \in \{1,\,2, \ldots, n\}$, the parameters corresponding to the first $k$ sets of bijective transformations are estimated sequentially, in some sense of optimality, for constructing flexible probability distributions, regardless of the remaining $(n-k)$ sets of bijective transformations. Moreover, our objects of interest are to highlight two complementary mathematical arguments that leverage an informatics-theoretic formalization, based-on empirical likelihood
estimators under moment restrictions, and a sequential decision-making paradigm, with mixing distributions, for updating and aggregating the estimated parameters of the overall normalizing flows. As a by-product, the framework provides a reliable surrogate model, conditioned on the model parameters defining the forward computational simulation, that allows samples generation, with statistical powers, and facilitates computationally tractable scheme in the
Bayesian paradigm for inference, hypothesis testings and uncertainty quantification.
Read the full abstract ▾
Machine Learning
Inference and learning in sparse autoencoders as natural gradient flow
Hadi Vafaii, Tejas Rao, David Chanin et al. (+6 more)
Sparse autoencoders are widely used to uncover interpretable features in neural networks, yet reliable recovery remains difficult when features overlap or activate infrequently. These challenges involve both inferring which features explain an input and learning the dictionary that represents them. Here, we unify inference and dictionary learning as natural-gradient flows on a shared variational free energy. We instantiate this framework as BeFOND, an encoder-free sparse coding model with closed-form inference and learning dynamics. We show how recurrent explaining away reduces interference between overlapping features, while Fisher preconditioning can compensate for the slow learning of rare features. On synthetic data, BeFOND improves dictionary recovery and rare-feature detection, with a growing advantage over amortized baselines as superposition increases. On
language-model activations, it improves single-feature concept detection and selective intervention, outperforming pretrained reference SAEs with substantially less training data. Its feature quality continues to improve with dictionary width, whereas the evaluated baselines largely plateau. Together, these results show how improving inference and learning within a unified probabilistic framework can make better use of data and dictionary capacity to interpret and intervene on neural representations.
Read the full abstract ▾
All authors Hadi Vafaii, Tejas Rao, David Chanin, Thomas Fel, Jacob L. Yates, Bruno Olshausen, David Klindt, Dileep George, Miguel Lázaro-Gredilla
Machine Learning · Cross-listed from stat.AP
DeepAJM: Deep Association Joint Model for Irregularly Sampled data
Barsha Halder and Jeffrey A. Thompson
Joint Models simultaneously model longitudinal and survival outcomes, leveraging patterns in patients' longitudinal trajectory to improve the prediction of survival outcomes. The classical parametric joint models, however, rely on fixed parametric assumptions, making them susceptible to bias under model misspecification and smaller sample sizes. We propose a deep joint model, DeepAJM, that does not require any parametric assumptions, while retaining a partially interpretable, per-longitudinal-outcome association structure. The joint model uses an encoder-decoder (sequence-to-sequence) architecture to learn the latent structure in patients' time-varying
covariate trajectories. The model links the longitudinal processes to the survival processes through a learned interpretable association structure, in which each longitudinal output from the decoder gets remodulated by baseline
covariates before it contributes to the risk scores from the survival head of the architecture. The model was evaluated on three
datasets ( a cardiovascular-disease EHR cohort, a primary biliary cirrhosis (PBC2)
dataset, and a simulated
dataset) against a classical parametric joint model, TransformerJM, DA-LSTM and a Cox-based survival-only model. All models were assessed using C-index, integrated brier score (IBS), time-dependent AUROC, and time-dependent AUPRC. Our model achieved the best discrimination in terms of the C-index, time-dependent AUROC, and AUPRC across all
datasets.
Read the full abstract ▾
Machine Learning · Cross-listed from stat.ML
HyperNSDE: Personalized Neural SDEs for Joint Static-Longitudinal Clinical Data Generation
Perrine Chassat and Agathe Guilloux
Synthetic patient data generation is a promising solution to the dual challenge of data scarcity and
privacy constraints in healthcare machine learning. Realistic synthesis of patient-level clinical data requires jointly modeling heterogeneous static
covariates, irregularly sampled longitudinal trajectories, and informative observation times - three tightly coupled components in practice yet rarely addressed together. We propose HyperNSDE, a continuous-time generative model that conditions a latent Neural SDE on static patient representations through a hypernetwork, allowing baseline characteristics to shape trajectory
evolution beyond the initial condition without requiring a trajectory encoder, while
stochastic latent dynamics capture realistic variability in generated paths. Observation times are modeled jointly through a latent-state-dependent intensity process, and training on irregular
stochastic paths is stabilized via a deterministic-
stochastic path decomposition with a non-
adversarial signature-kernel objective. Experiments on simulated and real clinical
datasets show improved observation-time fidelity and competitive performance, while matched-grid analyses reveal that forecasting and correlation metrics are affected by observation-grid regularity and trajectory smoothness.
Read the full abstract ▾
Machine Learning
Multigroup Fairness and Omniprediction: Separations and Equivalences
Sílvia Casacuberta, Parikshit Gopalan, Varun Kanade et al. (+3 more)
Omniprediction is a learning guarantee which requires a single predictor to be competitive relative to the best hypothesis from a
benchmark class for any loss chosen from a family of loss functions. Loss Outcome Indistinguishability (loss OI for short) is a stronger notion that implies omniprediction. It requires the predicted distribution on labels to be indistinguishable from the true distribution to tests that depend on the loss functions and the
benchmark class. Multiaccuracy and multicalibration are multigroup
fairness notions that generalize classical notions of calibration and accuracy in expectation. Most known learning
algorithms for omniprediction (both for the standard notion and for strengthenings like loss OI) rely on some version of these multigroup
fairness notions, or on an intermediate notion called calibrated multiaccuracy. We ask if this is necessary: Does omniprediction require some form of multigroup
fairness? We show that the answer is no for (plain) omniprediction, and yes for loss OI. First, a sequence of works shows that multicalibration or calibrated multiaccuracy imply omniprediction. We rule out even a weak converse, by showing that omniprediction for proper losses does not imply even accuracy in expectation, a much weaker notion than any of calibration, multiaccuracy, or multicalibration. Second, prior work showed how to achieve loss OI from a combination of calibration and multiaccuracy. We show a converse: loss OI is equivalent to a form of calibrated multiaccuracy.
Read the full abstract ▾
All authors Sílvia Casacuberta, Parikshit Gopalan, Varun Kanade, Omer Reingold, Konstantinos Stavropoulos, Pranay Tankala
Machine Learning · Cross-listed from cs.CV
Identity-Conditioned Score Fusion for Open-Set Person Re-Identification
Manyi Yao, Jurijs Nazarovs, Eunji Chong et al. (+6 more)
Robust person re-identification often combines complementary cues such as face, gait, and body shape. While adaptive fusion typically targets query quality, model strength also varies across identities. We introduce identity-conditioned score fusion, a framework that tailors weights to each gallery identity without training. By contrasting intra-identity consistency against cross-identity impostors, it extracts identity-specific profiles that couple with query-conditioned adaptation via a parameter-free rule. This widens the separation between true and false matches while preserving score calibration. Evaluations on three clothes-changing person re-identification
benchmarks show that our method consistently outperforms statistical, rank-based, and learned baselines, achieving up to an 8.8% absolute reduction in the false non-identification rate and demonstrating the value of identity-conditioned fusion in open-set person re-identification.
Read the full abstract ▾
All authors Manyi Yao, Jurijs Nazarovs, Eunji Chong, Abhishek Sharma, Rohan Sarkar, Yue Guo, Christian R. Shelton, Amit K. Roy-Chowdhury, Debashish Pal
Machine Learning
Dynamic Budget Allocation for LLM Evaluation under Hard Resource Constraints
Shai Feldman and Yaniv Romano
We evaluate large
language models (LLMs) in multi-turn interactions through their time-to-event: the number of interaction steps required to produce an event of interest, such as a successful jailbreak or agentic task completion. Under limited compute, interactions may be terminated before the event occurs, so that event times are only partially observed (censored). Existing allocation methods for calibrating time-to-event bounds satisfy the budget only in expectation and can exceed the available budget on a particular evaluation run. Enforcing a hard constraint is particularly challenging as the cost of a trajectory is initially unknown. We introduce Hard-budget Allocation with Reflow for Predictive calibration (HARP), a budget allocation that satisfies hard resource constraints and adaptively reallocates unused budget. We show how to use HARP to construct lower predictive bounds (LPBs) on the time-to-event and to estimate evaluation metrics such as the jailbreak rate on a fixed
benchmark. Although HARP induces dependence in acquisition decisions across different trajectories, we prove that HARP never exceeds the target budget, that its LPBs have finite-sample coverage guarantees, and that its metric estimates are unbiased. Experiments on agentic task success, LLM jailbreaks, toxic content generation, and RAG
hallucinations show that HARP achieves coverage close to the nominal level with low variance, while never exceeding the given budget.
Read the full abstract ▾
Machine Learning
Towards Explainable Benchmarking for Data-driven Post-Wildfire Debris Flow Prediction
Zhisheng Qi, Li Zhu, Utkarsh Sahu et al. (+3 more)
Post-wildfire debris flows (PFDFs) are destructive sediment-laden hazards triggered when intense rainfall strikes recently burned terrain, destabilizing hillslopes and threatening infrastructure, local economies, and community safety. Data-driven methods have been proposed to learn predictive patterns directly from historical PFDF observations. However, the current research landscape of data-driven PFDF prediction remains highly fragmented across feature spaces, model architectures, and evaluation protocols, making rigorous comparison and the derivation of scientific insights difficult. Moreover, existing studies lack a systematic investigation into the relative importance of heterogeneous factors (e.g., meteorological conditions, terrain characteristics, soil properties, and burn severity) in triggering PFDF. To address these limitations, we present a unified
benchmark for data-driven PFDF prediction, enabling fair and comprehensive evaluation across diverse models and feature configurations. Furthermore, to better understand the underlying drivers of PFDF formation, we propose a reinforcement learning-based feature selection framework that identifies factors whose perturbations render positive and negative events indistinguishable, thereby discovering the regional underlying mechanisms of PFDF occurrence across regions. Our code and
benchmark are publicly available at https://github.com/KINDLab-Fly/PFDF-
Benchmark.
Read the full abstract ▾
All authors Zhisheng Qi, Li Zhu, Utkarsh Sahu, Douglas Tommey, Josh Roering, Yu Wang
Machine Learning
RELACE: retrospective likelihood-based action credit estimation for long-horizon language agents
Sayak Chakrabarti and Sathish Reddy Indurthi
Group Relative Policy Optimization (GRPO) avoids a separate critic by estimating advantages from rollout groups. For multi-turn
agents, however, trajectory-level supervision provides coarse, noisy credit: terminal rewards do not locate errors and can penalize useful actions alongside mistakes. Group-in-Group Policy Optimization (GiGPO) and subsequent methods refine supervision through state-conditioned comparisons, but their credit estimates remain sensitive to downstream decisions and outcomes. We introduce RELACE, Retrospective Likelihood-based Action, a critic-free framework that integrates retrospective action assessment with state-conditioned advantage estimation. RELACE evaluates executed actions through teacher-forced likelihood scoring under both their original contexts and outcome-augmented contexts. Comparing these likelihoods yields a trajectory-normalized retrospective factor that captures outcome-dependent changes in action plausibility, rather than hindsight plausibility alone. We use this factor to reweight discounted task returns and construct local advantages by comparing weighted returns among actions from equivalent states within a task. This couples retrospective relevance with observed reward, producing fine-grained credit that complements trajectory-level GRPO supervision. Temporal smoothing and success-protecting masking further stabilize the local signal. RELACE requires neither auxiliary value nor reward models nor additional autoregressive rollouts for credit estimation. Experiments on ALFWorld and WebShop with Qwen2.5-1.5B-Instruct and Qwen2.5-7B-Instruct demonstrate substantial improvements over GRPO, GiGPO, and HCAPO. With the 1.5B model, RELACE achieves $96.35\%$ success on ALFWorld and $79.43\%$ on WebShop, surpassing GiGPO by $5.47$ and $5.60$ percentage points, respectively.
Read the full abstract ▾
Machine Learning
Stepped MoE: Segment-Level Routing with Configurable Inference Complexity
Arnav Kundu, Zhaoyang Xu, Bairu Hou et al. (+3 more)
Training large
language models (LLMs) is resource-intensive, and adapting them for diverse deployment scenarios with varying computational constraints remains challenging. While elastic architectures enable flexible model deployment and sparsely activated models allow input-adaptive computation, existing approaches treat these dimensions independently. Moreover, models catered towards on-device edge inference need to conform to the
memory and compute limitations of the serving devices. In this paper, we introduce a unified framework that combines elastic structures with sparsely gated architectures to create models that adapt simultaneously to both deployment constraints and task requirements. Our approach employs a model backbone that conditions on both the context and target efficiency specifications, enabling fine-grained control over the accuracy-efficiency trade-off at inference time. The model learns to activate task-relevant parameters within elastically-nested sub-networks, allowing a single model to span multiple capacity points while maintaining input-adaptive routing. Through experiments we demonstrate that we can create a model that allows the flexibility to use 1,2,3,4 billion parameters while being more accurate than their dense counter-parts (2-5% on knowledge-intensive
benchmarks) and at par with their static versions while delivering similar latency metrics as dense models. Overall, we save on device disk space by sharing the model parameters, allow flexibility of serving based on DRAM and compute available while delivering more accurate results.
Read the full abstract ▾
All authors Arnav Kundu, Zhaoyang Xu, Bairu Hou, Chang Gao, Reed Li, Tao Lei
Machine Learning · Cross-listed from cs.CR
Evaluating Behavioral Context for Interpretable IAM Policy Risk Scoring in Cloud Environments
Yassin Elsharkawy
IAM policy analysis typically emphasizes the authorization capabilities encoded in a policy, but security analyst review priority may also depend on the behavioral and environmental context surrounding a policy event. This paper evaluates whether contextual information provides measurable incremental value for interpretable IAM policy risk prioritization beyond policy and effective-authorization information. AWS is used as the experimental cloud provider because its IAM and audit-telemetry ecosystem enables controlled evaluation using AWS IAM Context Bench, a
benchmark containing 534 real AWS experimental observations across policy, environment, and behavioral scenarios, including matched cases where policy and environment remain fixed while behavioral context changes. Three Explainable Boosting Machine models are evaluated under the same leakage-controlled grouped cross-validation protocol: a policy-centric baseline, a policy-plus-environment model, and a full-context model incorporating CloudTrail telemetry. The full-context model substantially reduces analyst-priority prediction error relative to the policy-centric baseline and closely tracks the reference priority ordering. In matched same-policy context pairs, the policy-centric model remains invariant, whereas the full-context model separates benign and suspicious behavioral conditions with high directional accuracy. The results also show improved concentration of high-priority cases at the top of simulated analyst review queues. These findings indicate that behavioral and environmental context can provide useful incremental information for analyst-oriented IAM risk prioritization while preserving an interpretable additive model structure. The formulation is applicable beyond AWS conceptually, although cross-provider validation remains future work.
Read the full abstract ▾
Machine Learning
CausalBind: Causal Modeling and Learning for Protein-Molecule Virtual Screening
Loka Li, Jin Tian and Kun Zhang
Protein-molecule virtual screening is increasingly cast as a problem of representation learning in a shared
embedding space. Existing methods rely on dense holistic
alignment, entangling invariant binding determinants with nuisance correlations and limiting transfer to new targets. It has been noted that binding in
protein-molecule systems involves sparse cross-modality interactions: binding is governed by a small contact interface and a few decisive local interactions (e.g., hydrogen bonds, hydrophobic contacts, and salt bridges) rather than the global structures of the
protein and molecule. We hypothesize that uncovering and leveraging sparse interaction patterns is critical for generalization beyond the training data, as these patterns are reusable and expected to improve performance across different scenarios. In this paper, we aim to identify and leverage sparse interaction patterns, and verify our hypothesis. Since the training data contain only observed binding pairs, we formalize this prior via a V-structure causal model under Heckman-style selection, and establish three theoretical results: (i) the latent concepts of interacting
proteins and molecules are not identifiable without appropriate sparsity constraints; (ii) these concepts and their sparse interactions are component-wise identifiable under structural sparsity conditions; and (iii) a low-rank relaxation of these conditions yields subspace identifiability of the concepts and interactions. Inspired by these principles, we propose CausalBind with three implementation variants. Extensive experiments on DUD-E and LIT-PCBA
benchmarks show that all variants consistently outperform strong
retrieval baselines, with the largest gains on LIT-PCBA early enrichment, and further generalize to target- and scaffold-level out-of-distribution splits. Code is available at https://github.com/lokali/CausalBind.
Read the full abstract ▾
Machine Learning · Cross-listed from eess.AS
Logbook: Extremely Long-form Audio Event Understanding
Kwanghee Choi, Suwon Shon, Dmitriy Serdyuk et al. (+7 more)
Audio
benchmarks are built around short, pre-segmented clips, limiting model design to brief inputs or fixed vocabularies. To close this gap, we introduce Logbook, a
benchmark for hour-scale audio understanding, with recordings ranging from ten minutes to six days. Given a continuous audio recording and an event label vocabulary, a system must predict a gap-free segmentation with an event label and a description per segment. We compare 52 systems, end-to-end and cascaded, and ablate fine-tuning, context length, and
reasoning budget. We find the task tractable, though the best systems remain below the human reference. Also, over-segmentation is pervasive, and fine-tuning partially mitigates it. Finally, end-to-end are often better than cascaded systems, but degrades with longer context.
Read the full abstract ▾
All authors Kwanghee Choi, Suwon Shon, Dmitriy Serdyuk, Guitang Lan, Chao-Wei Huang, Mohammad Sadegh Rasooli, Sangeeta Srivastava, Zhaojiang Lin, Saurabh Adya, Ming Sun
Machine Learning
Selective Critique for Cost-Aware LLM Agents in Long-Horizon Decision Making
Heewon Park, Somin Im and Minhae Kwon
Improving the reliability of large
language model (LLM)
agents in long-horizon decision-making remains a key challenge. When deployed as autonomous
agents interacting with complex environments, early mistakes can propagate through trajectories and cause cascading failures. Recent approaches improve reliability by incorporating external critique or deliberation, but invoking these mechanisms at every step substantially increases
token consumption and latency, limiting practical deployment. We propose SAG (Self-improving
Agent with Gated critique), a cost-aware framework that formulates critique invocation as a step-wise decision problem during long-horizon interaction. SAG introduces a lightweight, training-free gating mechanism that estimates the utility of critique using action-level ambiguity signals–global
entropy and local top-2 margin–computed over admissible actions. From a decision-theoretic perspective, this mechanism approximates the Value of Information (VoI) of critique, enabling the
agent to selectively allocate expensive feedback only when its expected benefit justifies the cost. SAG further incorporates online bootstrapped self-improvement, allowing the actor to internalize critic-assisted behaviors and progressively reduce reliance on critique. Across three long-horizon interactive
benchmarks and multiple backbone models, SAG substantially improves the performance-cost trade-off compared with both no-critique and always-on critique
agents. On ALFWorld, SAG increases task success from 24.6% to 78.4% while maintaining a
token budget comparable to ReAct, yielding a $3.1\times$ improvement in normalized
token efficiency. Moreover, a 7B actor with a lightweight 3B critic achieves performance comparable to a 14B actor without critique, showing that selective critique can recover most of the reliability benefits of deliberation while dramatically reducing inference cost.
Read the full abstract ▾
Machine Learning
Weight Oracles: Reading Neural Network Weights with Language Models
Krishna Kabra, Constantin Venhoff and Christian Schroeder de Witt
Interpretability methods for neural networks are predominantly reactive: they analyse activations produced during specific forward passes, requiring known inputs to find hidden capabilities such as backdoors. We propose Weight Oracles, fine-tuned
language models that diagnose properties of a target network by reading its raw weights directly, without behavioural testing. We investigate this paradigm in two phases. Phase I establishes feasibility: through a staged curriculum and an external chain-of-computation that delegates parameter-free operations to deterministic code, an explainer LLM learns to simulate the forward pass of small
transformers from their weights, achieving 99% holdout accuracy on unseen targets. Phase II repurposes this infrastructure for safety auditing. We train an oracle on natural
language diagnostic questions about weight anomalies using only benign pathologies as training signal, and evaluate it zero-shot on backdoors absent from training. The oracle achieves AUROC 0.93 on
attention-routed backdoors and 0.81 across a diversified threat distribution including stealth and adversarially regularized variants. Hand-crafted statistical detectors are sharp on the threat models they implicitly target but collapse on threat-model shift, while the oracle remains uniformly competent across attack types. Scaling to realistic model sizes remains the principal open challenge.
Read the full abstract ▾
Machine Learning
Structuring MoE Expert Selection for Agentic Reinforcement Learning
Bolian Li, Ting-Yao Hu, Cheng-Yu Hsieh et al. (+3 more)
Long-horizon LLM
agents are frequently implemented using sparse mixture-of-experts (MoE) models, yet the co-design of agentic behavior and MoE structures remains underexplored. In this work, we comprehensively study the connections between agentic post-training and MoE expert selection. In off-the-shelf MoE models, we observe expert selection exhibits a specialized structure that naturally aligns with agentic trajectories. Specifically, expert routing overlaps more between turns where the
agent performs semantically similar operations (e.g., READ, UPDATE) than between turns with differing operations. However, standard RL
algorithms ignore this specialization, allowing the MoE routing to go uncontrolled during training, which empirically limit task performance and inference efficiency. To address this, we introduce a hierarchical routing control framework for agentic tasks. We explicitly encourage turn-level expert selections to align with agentic operations while regularizing
token-level expert selections to maintain local consistency. To resolve stability issues that arise during post-training with the proposed methods, we further introduce an
entropy-gated control mechanism. Overall, our routing control framework achieves over 10-point improvements in success rate on all evaluated
benchmarks. These results demonstrate that agentic trajectory structure provides an effective signal for optimizing MoE capacity during RL post-training.
Read the full abstract ▾
All authors Bolian Li, Ting-Yao Hu, Cheng-Yu Hsieh, Sanjoy Chowdhury, Oncel Tuzel, Raviteja Vemulapalli
Machine Learning · Cross-listed from quant-ph
Advantage of Entangled Learning Rules in Quantum Measurement Class Learning
Arka Prabha Das and Abram Magner
Learning with data in the form of
quantum states is of current interest and has led to a variety of problems that boil down to interaction with the available data via
quantum measurement and classical post-processing of observed classical outcomes. In
quantum measurement PAC learning, one is given a sequence of unknown, prepared
quantum states and classical labels, along with a hypothesis class of candidate measurements. The task is to select a measurement from the hypothesis class that minimizes a fixed notion of error in prediction of the classical labels via measurement of a new state by the selected hypothesis. In this work, we consider the advantage of interacting with the given data in the measurement learning framework using learning rules given by measurements that cannot be implemented using local operations and classical communication (LOCC), as opposed to single-copy learning rules. We provide a construction showing that there exist learning scenarios wherein single-copy learning rules are asymptotically suboptimal compared to optimal ones. We then show that learning rules based on entangled measurements enjoy at most a
polynomial sample complexity advantage over single-copy learning rules in the PAC learning setting (under a natural joint measurability covering assumption).
Read the full abstract ▾
Machine Learning
Scale-Invariant Training for Time Series Foundation Models
Ignacy Stepka, Willa Potosnak, Kin G. Olivares et al. (+1 more)
Time series foundation models (TSFMs) are trained on large collections of time series
datasets that span various morphologies and domains. This setting exposes models to series whose scales – typical magnitudes of their values – can differ substantially. Affine scaling methods such as Reversible Instance Normalization (ReVIN) scale model inputs and reverse the transform before computing the loss. We show that this inversion multiplies each series' gradient by $b^p$ relative to loss on scaled targets, where $b$ is the scaling denominator (e.g., standard deviation) and $p$ is the loss degree. We call this scale-contaminated training (ScaleCon), because the scale of each series consequently becomes an importance weight, causing high-scale series to dominate training. For any scale-equivariant scaler and residual loss that is homogeneous of degree $p$, including MSE, MAE, and
Quantile Loss, we prove that computing loss on scaled targets makes every mini-batch gradient and, consequently, the full optimization trajectory invariant to arbitrary independent rescaling of the training series, yielding scale-invariant training (ScaleIn). Notably, existing TSFMs use both objectives, with neither consistent reporting nor a common convention on how to compute training loss. We isolate the convergence disparity induced by ScaleCon and its correction under ScaleIn in controlled studies on synthetic and real data. In pretraining across four TSFM architectures, ScaleIn lowers MASE in all 24 architecture-
benchmark comparisons, with average reductions across TSFMs of 18.8% on GIFT-Eval and 21.9% on the M-competitions. The gains extend to supervised neural forecasting, where it lowers MASE in 16 of 20 matched settings. Most existing time series forecasting pipelines can adopt ScaleIn with a one-line code change.
Read the full abstract ▾
All authors Ignacy Stepka, Willa Potosnak, Kin G. Olivares, Artur Dubrawski
Machine Learning
ATLAS-AL: Adaptive Trust-Region for Latent Adversarial Searches via Active Learning
Marsalis Gibson, Claire Tomlin and Shankar Sastry
Security evaluation of learning-based systems requires more than just testing the system against a fixed collection of attacks. It requires adaptive mechanisms that can efficiently discover sets of inputs that induce model failure. We introduce ATLAS (Adaptive Trust-Regions for Latent
Adversarial Searches), which is a query-based framework that discovers
adversarial input sets for black-box learning systems. ATLAS casts attack generation as an active learning level set estimation problem then combines calibrated approximations with a local-global sampling architecture to find regions of the input space that contain
adversarial examples. Once discovered, ATLAS is designed to sample points within these
adversarial regions to build
adversarial sets that accurately represent the state of robustness of the target model. When applied on toy experiments, we find that ATLAS is able to recover more of the
adversarial region under a limited query budget than does previous work. When applied to standard and adversarially trained MNIST, CIFAR, and ImageNet model targets, ATLAS produces better representative attacks than other query-based black-box attacks (NES, SignHunter, BayesOpt). ATLAS represents an automated red-teaming framework that can be used for both analyzing the robustness of learning-based systems under development and continuous auditing to see how the robustness of a system changes over time.
Read the full abstract ▾
Machine Learning · Cross-listed from stat.ML
Assumption-lean logistic regression with missing covariates
Jyotishka Ray Choudhury, Kabir Aladin Verchand, Richard J. Samworth et al. (+1 more)
Missing
covariates are frequently encountered in supervised learning problems, and classical methods for estimation using such data use carefully chosen imputation schemes for missing data, or likelihood approximations that lead to nonconvex $M$-estimation problems. These methods and their relatives are suitable for scenarios in which the
covariate distribution is known, and more broadly, have enjoyed tremendous success in linear models. But even in basic nonlinear problems such as logistic regression in moderate dimensions, such methods can experience drastic failure modes when the
covariate distribution is unknown. Motivated by the need for reliable alternatives, we consider the problem of parameter estimation in logistic regression with missing
covariates. Crucially, we operate in the assumption-lean setting where the
covariate distribution is unknown (but bounded). We design a
stochastic approximation method that is based on $Z$-estimation with a novel monotone operator, and establish that our
algorithm is computationally efficient and achieves provable signal recovery at parametric rates under the hypothesis that
covariates are missing completely at random. Our theory sharply characterizes the $\ell_2^2$ risk of the
estimator in terms of the missingness profile, accommodating heterogeneous observation probabilities. Importantly, it shows that our method always outperforms the de facto “complete-case”
estimator that ignores observations with any missing data. Even in the setting with homogeneous missingness (in which each
covariate is observed independently with probability $q$), our bounds exhibit intricate and nonstandard dependence on $q$ that can yield significant improvements over using only complete cases. We complement our upper bounds with new information-theoretic lower bounds that show that this intricate dependence on $q$ is fundamental in a minimax sense.
Read the full abstract ▾
All authors Jyotishka Ray Choudhury, Kabir Aladin Verchand, Richard J. Samworth, Ashwin Pananjady
Machine Learning · Cross-listed from math.OC
A Single-Loop, Constant-Batch First-Order Penalty Method for Stochastic Bilevel Optimization
Xingyu Chen, Ming Yang, Quanqi Hu et al. (+1 more)
Recent advances in penalty-based methods for
stochastic bilevel optimization (SBO) have eliminated the need for second-order derivative oracles. However, for
stochastic nonconvex-strongly convex bilevel problems, existing first-order methods typically rely on nested loops and/or large batch sizes for attaining $O(\epsilon^{-6})$ or $O(\epsilon^{-4})$ sample complexity under standard bounded-variance assumption or mean-square smoothness assumption. Achieving these rates with a single-loop penalty method and a constant batch size remains challenging due to a large penalty value needed for an accurate approximation. To address this challenge, we develop a
stochastic SIngle-loop COnstant-Batch first-order penalty method (SICO) that combines two complementary ingredients. First, it performs one
stochastic-gradient update per-iteration for both the original lower-level and penalized problems, with a projection that controls the separation between their iterates. Second, it applies an exponential moving average to stabilize the upper-level gradient
estimator. We show that this combination achieves $ O(\epsilon^{-6}) $ sample complexity using only $O(1)$
stochastic-gradient samples per iteration under unbiased, bounded-variance
stochastic gradients. Under the additional mean-square smoothness assumption on the lower-level
stochastic gradients, the same
algorithm improves the complexity to $O(\epsilon^{-4})$ also with $O(1)$ batch size. To the best of our knowledge, this is the first work to match the best-known convergence rate for fully first-order SBO methods using a single loop and a constant batch size. This result addresses an open problem posed in the literature.
Read the full abstract ▾
All authors Xingyu Chen, Ming Yang, Quanqi Hu, Tianbao Yang
Machine Learning
FlexiFlow: Bandit-based Model Switching in ML Workflows
Abhilash Jindal, Todd Nief, Bhanu Prakash Vangala et al. (+6 more)
Model optimizations help improve inference performance and accuracy of ML workflows. However, relying on a single model to perform inference across all data batches often fails to maximize accuracy and thus overall performance. In many cases, alternate models could perform better on specific subsets of data where a primary model underperforms. Our experiments with real ML workflows indeed show that switching models improves workflow accuracy by up to 23%. Yet, current systems lack the ability to adaptively switch between models based on performance, forcing users to manually test models in sequence. We present FlexiFlow, a dataflow system that dynamically switches between alternate models when the current model exhibits low accuracy. FlexiFlow learns to rank models using a novel multi-armed bandit approach that accounts for model runtimes, probability of passing user-defined assertions, and the computational structure of the ML workflow. We show that the standard Thompson sampling approach is insufficient for switching models in ML workflows. In contrast, our proposed approaches are effective and scales to complex real-world ML workflows. Experiments show that switching models at runtime while reusing intermediate results provides higher accuracy, but also 48% efficiency gain compared to sequential workflow runs.
Read the full abstract ▾
All authors Abhilash Jindal, Todd Nief, Bhanu Prakash Vangala, Shankaradithyaa V, Tvisha Malik, Anshik Sahu, Aaron Schein, Amitabh Chaudhary, Tanu Malik
Machine Learning
Lock-in EP: An In-Situ Training Algorithm for Oscillatory Hardware
Sowjanya Tammali and Wilkie Olin-Ammentorp
Analog hardware platforms offer the potential to reduce energy consumption over digital architectures, but in order to succeed, large-scale analog systems must also be able to operate with or recover from the variability of their components. Towards this goal, we derive and demonstrate the lock-in
equilibrium propagation (LIEP) training method. LIEP provides local gradient information for each component in an oscillatory network without separate forward and backward sweeps, potentially allowing for in-situ learning capabilities on analog oscillatory hardware platforms. We demonstrate that LIEP can be used both for ab-initio training as well as recovering performance when pre-trained parameters are perturbed. We show that LIEP can be formulated as a three-factor update rule, and suggest that although the method is currently only validated on shallow networks, alternate architectures may allow it to extend to deep and large-scale networks addressing complex tasks.
Read the full abstract ▾
Machine Learning
Algorithmically Aligned Neural Agglomerative Tree Construction
Robert R Nerem, Pranav Singh, Cheyenne Ward et al. (+1 more)
Linkage
algorithms for hierarchical clustering (HC) are a powerful and efficient framework for constructing clustering trees, yet it is often unclear which merge rule best suits a given
dataset or task. In contrast, neural approaches can learn from data, but often fail to retain the efficiency and size generalization of classical
algorithms. We introduce NN-linkage, a neural network (NN) model that can learn task-specific and locally dependent merge rules while retaining the recursive structure and efficient inference of classical linkage
algorithms. In particular, our model is algorithmically aligned with the Lance-Williams (LW) recurrence, a parameterized framework for defining a broad, continuous family of linkage rules for agglomerative HC. Classical methods such as single linkage (SL), complete linkage (CL), and average linkage arise as discrete choices within this broader family. We show that NN-linkage is a universal approximator for continuous linkage functions, including LW recurrences, and, when paired with a
transformer encoding, can also approximate globally dependent rules such as robust single-linkage. We further show that NN-linkage can exactly implement any symmetric constant-coefficient LW recurrence across all input sizes. On the empirical front, we evaluate NN-linkage in real-world applications, clock-tree routing and phylogenetic reconstruction, using both synthetic and real
datasets, demonstrating its effectiveness over both classical
algorithms and other neural approaches. By learning merge rules directly from target trees, NN-linkage extends efficient HC to scientific and engineering objectives not adequately captured by existing hand-designed linkage rules.
Read the full abstract ▾
All authors Robert R Nerem, Pranav Singh, Cheyenne Ward, Yusu Wang
Machine Learning · Cross-listed from cs.CV
What Words Keep of a Place: Zero-Shot Language Reasoning for Cross-View Geo-Localization
Ayesh Abu Lehyeh, Jay Hwasung Jung and Safwan Wshah
Cross-view geo-localization is commonly solved as an image
retrieval problem, matching a ground-level image against a database of satellite tiles through a jointly trained
embedding. Such models are accurate, but they need large paired supervision and cannot show what evidence supports a match. In this paper, we study a different question: how much of this task can be solved through
language alone? We
prompt a
multimodal large
language model (MLLM) to describe each ground panorama and each satellite tile as structured text, and localize by comparing these descriptions. No component is trained. We evaluate on 9,826 VIGOR pairs from four U.S. cities, in three settings. First, the descriptions are faithful but not discriminative. They agree closely across the two views, yet ranking the full pool by description similarity almost never returns the correct tile (0.39% Recall@1). Second, we narrow the pool to ten neighboring tiles, as a coarse prior would do. The same descriptions now become useful: an MLLM judge that scores structural consistency doubles random ranking and matches a strong lexical baseline. It also states which fields of the two descriptions agree and which conflict, which an
embedding distance cannot do, and which we see as a step toward interpretable localization. Third, we place the judge on a trained visual retriever. On the queries it ranks wrongly, reranking from images works, while reranking from our descriptions does not (23.5% against 10.7% Recall@1). Scene structure survives the conversion into
language, while the fine appearance detail needed to separate nearby places does not. Code and
prompts are publicly available at https://github.com/AyeshAbuLehyeh/GeoLingual.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CR
Lineage-Aware Memory Governance: A Derivation-Gated Framework for Privacy-Preserving Column-Level Access Control in Enterprise AI Agents
Venkata M Sangaraju and Sudhir Vissa
Enterprise AI
agents that share a
memory store face two unaddressed risks: sensitive data can leak through legitimately computed results the requester could not derive, and departments can silently compute a same-named key performance indicator (KPI) through conflicting logic. Existing
agent-
memory systems (e.g., MemGPT, Zep, A-MEM) gate
retrieval by content, ownership, and role, not derivation, missing a cached insight that embeds a forbidden column. We introduce the Analytical
Memory Unit (AMU), a
memory schema that attaches a full derivation (lineage) graph to every cached result, gated by a
retrieval policy that serves a hit only when the requester is authorised for every column touched. Provided lineage recording is complete, we prove by construction that the policy blocks
retrieval of results derived from a sensitive column outside the requester's permissions, at O(n) worst case – a conditional design guarantee, not an empirical claim, that excludes derived features encoding sensitive information without naming their source. Eliminating measured leakage required 75-90% recorded lineage completeness, so we treat 90% as a conservative deployment target. Across six experiments, lineage-gated
retrieval removes the 18.8-25.5% cross-department leakage naive content-gated
memory suffers, keeping 81.5-82.6% of
memory reuse at 13.8 microsecond worst-case overhead. A real-
agent proof-of-concept with LLM-generated SQL is consistent with the guarantee: zero leaks over 9 round-trips, two conflicts caught automatically – though a feasibility demonstration, not evidence of production viability. This offers a practical governance layer for shared
agent memory, complementing source-layer access control and supporting EU AI Act compliance.
Read the full abstract ▾
Machine Learning
Neural Algorithmic Reasoning for Graph Saddle Point Problems
Samantha Chen, Jesse He, Coleman Clougherty et al. (+2 more)
Neural algorithmic
reasoning, or aligning a neural network with an algorithmic paradigm, has emerged as an approach to solving
polynomial-time-solvable and computationally harder combinatorial optimization problems. We propose a new message-passing framework based on the Chambolle-Pock Primal–Dual Hybrid Gradient (PDHG) method called GraphPDHG for solving general graph saddle-point problems. Theoretically, we show that GraphPDHG can efficiently solve a family of graph saddle-point problems by simulating PDHG. We also show that our network can learn an accelerated PDHG
algorithm. Experimentally, we support our results on accelerated PDHG by evaluating the performance of our model as a learned warm start for second-order optimization techniques (SSNAL). We also show that
alignment with PDHG leads to stronger size generalization than non-aligned graph neural network (GNN) baselines. Overall, we propose a novel architecture for solving a general family of optimization problems on graphs.
Read the full abstract ▾
All authors Samantha Chen, Jesse He, Coleman Clougherty, Gal Mishne, Chester Holtz