Machine Learning · Cross-listed from stat.ML
Estimating Mixture Distributions via Stochastic Mirror Descent
Mohammadreza Ahmadypour, Tara Javidi and Farinaz Koushanfar
We revisit the classical problem of estimating an unknown distribution from its samples by fitting a mixture model that minimizes cross-
entropy loss. Framing the task as a
stochastic convex optimization problem over the space of $ M $-component mixture distributions, we propose a family of
estimators derived from the
stochastic mirror descent (SMD)
algorithm. This optimization-based approach provides a principled and flexible framework that generalizes traditional
estimators and proposes a variety of novel
estimators through the choice of Bregman divergences. A key advantage of our method is that it scales efficiently with the number of candidate components $ f_i $; that is, one can employ a large set of basis distributions in the mixture model without incurring significant computational overhead. This enables richer approximations and improved estimation accuracy. Moreover, in the case of categorical distribution (discrete outcomes) our
estimators do not require a strict lower bound, in other words our framework does not require the precise knowledge of the support of the distribution. We demonstrate that, under mild conditions, the proposed $ \varphi $-SMD
estimators achieve near-optimal convergence rates in both Kullback-Leibler (KL) divergence and $ \ell_2 $-norm and offer practical benefits when computation is expensive. Our numerical analysis highlights improved performance guaranties over classical
estimators, particularly in terms of sample efficiency and scalability.
Read the full abstract ▾
Machine Learning
BandVQ: Band-Wise Vector-Quantized EEG Foundation Model
Jamiyan Sukhbaatar, Satoshi Imamura and Toshihisa Tanaka
A central challenge in electroencephalography (
EEG) foundation modeling is learning transferable representations across recordings with diverse tasks, montages, references, and spectral characteristics. Existing masked modeling approaches often rely on broadband continuous patches or a single discrete representation, which may underrepresent frequency-specific activity. This paper proposes BandVQ, a band-wise vector-quantized
EEG foundation model that decomposes
EEG into delta, theta, alpha, beta, and gamma bands, trains an independent VQ-VAE tokenizer for each band, and pretrains a shared
Transformer encoder on the resulting discrete VQ code indices. The encoder uses masked code
tokens, quantized absolute log-power
tokens, channel and temporal
embeddings, and metadata prefix
tokens representing reference, band, task family, and phase. Region-based masking is also introduced to reduce the trivial reconstruction of spatially adjacent electrodes. The model is pretrained on 71 public
EEG corpora comprising over 9,200 subjects and 357,000 single-channel hours and evaluated on six subject-independent classification
datasets. Under the current evaluation setting, the proposed model achieves strong transfer performance, with the highest reported results on three cognitive tasks and competitive performance on three motor imagery tasks.
Read the full abstract ▾
Machine Learning
Quaternion Self-Attention with Shared Scores
Shogo Yamauchi, Tohru Nitta and Hideaki Tamori
Quaternion neural networks are parameter-efficient and model multidimensional dependencies by representing four related features as a single entity. However, existing quaternion self-
attention computes component-wise scores and applies independent softmax operations to each component, which increases the computational cost and allows
attention distributions to diverge across components. We propose a shared-score quaternion self-
attention mechanism that computes a single real-valued score using the quaternion inner product and applies a shared
attention distribution across all components. This reduces score-computation multiplications by 75% and the number of softmax operations from four to one. We prove that, when queries and keys are produced by quaternion linear projections that induce component pre-mixing, the component-wise and shared scores lie in the same interaction subspace, indicating that independent component-wise
attention primarily re-parameterizes the same interactions rather than expanding the feature interaction space. In
speech enhancement, our method reduces inference time by up to 44.3% on a GPU and 58.1% on a CPU while maintaining quality, with consistent trends across
vision and natural
language processing.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.IR
MVR-cache: Optimizing Semantic Caching via Multi-Vector Retrieval and Learned Prompt Segmentation
Ali Noshad, Zishan Zheng and Yinjun Wu
To reduce LLM costs and latency, semantic caching systems must accurately identify when a new
prompt matches a cached one. Current methods often rely on simplistic similarity measures, which limit their effectiveness. We introduce MVR-cache, a novel semantic caching approach that significantly improves
retrieval accuracy by integrating Multi-Vector
Retrieval (MVR). MVR-cache is built upon a learnable
segmentation model that intelligently splits
prompts, enabling fine-grained similarity comparisons via MaxSim. We derive the model's training objective from a rigorous theoretical analysis. This can ensure that optimizing this objective directly maximizes cache hits under strict correctness constraints. To solve the resulting non-differentiable combinatorial optimization problem, we leverage a reinforcement learning-based training strategy with the theoretically grounded objectives as the reward. Experimental results on established
benchmarks across diverse tasks confirm that in comparison to the state-of-the-art, MVR-cache consistently increases the cache hit rates by up to 37% while maintaining the same correctness guarantees. MVR-cache is available at https://github.com/PKU-SDS-lab/MVR-Cache
Read the full abstract ▾
Machine Learning
Explainable Retinal Imaging for Prediction of Multi-Organ Dysfunction in Type 2 Diabetes
Mini Han Wang, Liting Huang, Wei Hong et al. (+1 more)
Background: Type 2 diabetes mellitus (T2DM) is increasingly recognised as a systemic disease characterised by coordinated dysfunction across metabolic, renal, lipid, and inflammatory pathways. Existing clinical assessments often fail to capture this multi-dimensional burden. Methods: We conducted a retrospective study of 1,195 patients using routinely collected laboratory biomarkers. System-level abnormality indices were constructed to quantify organ-specific dysfunction, and multi-system involvement was defined as abnormalities in two or more systems. Supervised machine learning models, including logistic regression, random forest, and gradient boosting, were trained to predict multi-system dysregulation. Model interpretability was achieved using SHapley Additive exPlanations (SHAP). Results: The gradient boosting model demonstrated near-perfect discrimination (AUC = 1.000), significantly outperforming logistic regression (AUC = 0.925). Feature attribution analysis revealed that hyperglycaemia, renal impairment, dyslipidaemia, and inflammation were the dominant drivers of multi-system risk. Dose-response relationships observed in partial dependence analyses further supported the biological plausibility of model predictions. Conclusion: This study presents an interpretable, data-driven framework for quantifying systemic disease burden in T2DM. By linking routine biomarkers to multi-organ dysfunction, our approach provides both predictive accuracy and mechanistic insight, offering potential for improved risk stratification and precision medicine in diabetes care. The data and code used in this study are openly available on GitHub at: https://github.com/MiniHanWang/Type-2-Diabetes-1.git
Read the full abstract ▾
All authors Mini Han Wang, Liting Huang, Wei Hong, Boonthawan Wingwon
Machine Learning
ROAR: Retrieval Opportunity-Aware Refinement for Zero-Shot Time Series Forecasting
Jinjin Chi, Lei Feng, Yiming Wang et al. (+4 more)
Retrieval augmentation provides time series forecasters with historical continuations, yet even candidates that outperform the base forecast may fail to improve the final prediction. We propose ROAR, a
Retrieval Opportunity-Aware Refinement framework for zero-shot time series
forecasting. To better exploit these improvement opportunities, its training objective allocates additional emphasis across queries based on base-forecast difficulty and the relative improvement offered by retrieved candidates. Using this objective, ROAR first learns to aggregate aligned historical candidates and uses a learned gate to control their correction strength against a fixed base forecaster. It then jointly calibrates the
forecasting module and gate to coordinate their contributions, while anchoring the combined prediction to the first-stage refined output. We derive exact decompositions of refinement gains and the opportunity-weighted training loss, and characterize additional gains from joint calibration under a local linearization. Experiments on seven
benchmarks show that ROAR achieves the lowest average MSE among the compared methods. Further evaluations demonstrate average
forecasting improvements across multiple backbone families and
retrieval-augmented forecasters.
Read the full abstract ▾
All authors Jinjin Chi, Lei Feng, Yiming Wang, Lulu Zhang, Yongcheng Jing, Chenyi Li, Ximing Li
Machine Learning
On the Impact of Class Imbalance on the Learning Dynamics of Deep Neural Networks:An Intuitive Insight
Ismail B. Mustapha, Shafaatunnur Hasan, Sunday O. Olatunji et al. (+1 more)
Class imbalance in deep neural networks (DNNs) has witnessed a rapid increase in research
attention in recent years. However, the varying accounts of the reasons behind the poor performance of DNN on imbalance data in pertinent literature shows that little is known about how this agelong phenomenon impacts the performance of DNNs. A better understanding of this problem is crucial to developing effective DNN-based imbalance methods. Thus, this study systematically investigates the impact of class imbalance on the learning dynamics of DNN by monitoring the learning pattern of DNN models on both the majority and minority classes of
datasets of varying imbalance ratios. Experimental findings shows that as against learning from balanced
datasets where DNN learns the classes similarly, class imbalance has severe deteriorating impact on the performance of DNN, driving the model to underfit the minority class samples in the early training epochs while simultaneously learning only the majority class. Although DNN ultimately learns the minority samples, learning in this manner only results in learnt minority representations that are non-generalizable at test phase because they are merely overfitted to keep the overall training loss as low as possible.
Read the full abstract ▾
All authors Ismail B. Mustapha, Shafaatunnur Hasan, Sunday O. Olatunji, Hatem S. Y. Nabus
Machine Learning · Cross-listed from cs.CR
SEED: Semi-supervised Continual MalwarE Detection for Tackling ConcEpt Drift on a BuDget
Suresh Kumar Amalapuram, Bikraj Shresta, Siva Ram murthy Chebiyam et al. (+2 more)
Machine learning based malware detectors become obsolete over time due to concept drift in benign and malware applications. Recent methods rely on fully labeled data and use hierarchical contrastive loss (HCL) with active learning to improve robustness against drift by exploiting semantic structure in malware representations. However, obtaining labeled data in the security domain is difficult. Under partially labeled settings, HCL suffers significant performance degradation in detecting unseen malware, especially on
datasets such as BODMAS where strong semantic structure may not exist. In this paper, we propose SEED, a semantic-structure-agnostic method for malware detection under limited supervision. SEED combines a tailored binary cross-
entropy objective with semi-supervised continual learning and active learning. For partially labeled seen tasks, unlabeled samples are projected into a representation space constructed from previously seen data using singular value decomposition, and paired with suitable labeled samples to encourage representation consistency. For unseen tasks with fully unlabeled data, uncertainty is quantified using cosine distance in representation space, and the most uncertain samples are selected for analyst labeling. We evaluate SEED on both Windows and Android malware
datasets. Using only 20% labeled data on seen tasks, SEED achieves average AUT improvements of 40% on BODMAS and 14% on AndroZoo for unseen malware detection compared to HCL* (the semi-supervised adaptation of HCL), while remaining competitive on APIGraph. Finally, we introduce a delayed buffer update strategy to reduce label noise propagation during replay and improve learning stability.
Read the full abstract ▾
All authors Suresh Kumar Amalapuram, Bikraj Shresta, Siva Ram murthy Chebiyam, Bheemarjuna Reddy Tamma, Sumohana S Channappayya
Machine Learning · Cross-listed from cs.CL
When Reasoning Hurts: Source-Aware Evaluation of Frontier LLMs for Clinical SOAP Note Generation
Faizan Faisal
Reasoning-enabled LLMs perform strongly on medical
reasoning benchmarks, but it remains unclear whether these gains transfer to structured clinical documentation; we investigate this question using SOAP note generation from clinical dialogue in a source-aware
benchmark spanning OMI Health, ACI-Bench, and PriMock57. We evaluate GPT-5.4, DeepSeek-V4-Flash, and Gemma-4-E4B in a controlled 2x2 design that independently toggles provider-native
reasoning and same-source
retrieval-augmented generation (RAG). Outputs are assessed using seven automatic metrics alongside two reference-aware LLM judges. Both evaluation approaches agree that a non-
reasoning GPT-5.4 configuration achieves the highest overall quality, while DeepSeek-V4-Flash performs best among
reasoning-enabled configurations. Enabling
reasoning significantly degrades GPT-5.4 performance across all three
datasets, whereas same-source RAG yields smaller, model-dependent improvements. Overall, the findings indicate that stronger
reasoning capability should not be assumed to improve fidelity-sensitive SOAP note generation without dedicated, task-specific evaluation.
Read the full abstract ▾
Machine Learning
Efficient DP-SGD for LLMs with Randomized Clipping
Enayat Ullah, Sai Aparna Aketi, Devansh Gupta et al. (+2 more)
Large
language models (LLMs) are trained on vast
datasets that may contain sensitive information. Differential
privacy (DP), the de facto standard for formal
privacy guarantees, provides a principled framework for training LLMs with provable
privacy protection. However, state-of-the-art DP training implementations rely on fast gradient clipping techniques with
memory overhead $O(B \min\{T^2, d^2\})$, where $B$ is the batch size, $T$ is the sequence length, and $d$ is the model width. This becomes prohibitive as both model size and context length grow. We propose DP-SGD-RC, a novel variant of DP-SGD with randomized clipping that reduces
memory and compute complexity. DP-SGD-RC leverages
stochastic trace estimation methods, specifically Hutchinson's
estimator[Hutchinson, 1989] and its improved variant, Hutch++[Meyer et al., 2021], to reduce the
memory footprint of per-sample gradient norm estimation. We provide a tight
privacy analysis showing that DP-SGD-RC achieves noise multipliers competitive with deterministic clipping. Experiments fine-tuning Llama~3.2-1B on long-context
benchmarks spanning classification, question answering, and summarization tasks demonstrate that DP-SGD-RC matches baseline utility while significantly reducing
memory and compute requirements.
Read the full abstract ▾
All authors Enayat Ullah, Sai Aparna Aketi, Devansh Gupta, Huanyu Zhang, Meisam Razaviyayn
Machine Learning · Cross-listed from math.NA
IV-Net: A neural network for elliptic PDEs with random and highly varying coefficients
Shan Zhong and George Biros
We introduce a novel neural operator architecture designed to approximate solutions of linear
elliptic partial differential equations with high-contrast, spatially varying coefficients. The network, termed the Iterated V-shaped Net (IV-Net), realizes a mapping from the input coefficients and righthand side to the corresponding solution field. The architecture of IV-Net is informed by, and closely resembles, a V-cycle multigrid solver. The IV-Net model is parameterized via convolutional layers defined in the physical domain. For coercive problems with highly heterogeneous coefficients, the proposed network exhibits superior performance relative to a proper orthogonal decomposition (POD) approach and several existing neural operator architectures. For low-frequency oscillatory Helmholtz problems with smooth coefficients, its performance is similar to that of a
Fourier neural operator. We analyze the approximation error and convergence behavior of IV-Net, its data efficiency, and its dependence on the underlying discretization mesh. Furthermore, we demonstrate the practical effectiveness of the architecture through a series of numerical experiments, including applications to uncertainty quantification, inverse problems, and prediction of quantities of interest.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
Towards a Universal Causal Reasoner
Qirun Dai, Xiao Liu, Jiawei Zhang et al. (+3 more)
Despite the importance of causal
reasoning, training LLMs to reason causally remains underexplored. Existing data efforts mostly focus on benchmarking LLMs on specific aspects of causality, making them less suitable for training generalizable causal reasoners. To address this, we propose UniCo, a data generation framework that both (1) addresses 18 causal query types across Pearl's Causal Ladder and (2) translates natively symbolic examples into code and natural
language forms to simulate real-world use cases where causal terms are not explicitly specified. To ensure data quality, UniCo grounds answers with exact causal inference and filters cases with
reasoning shortcuts. Upon supervised finetuning with 66.6K UniCo-generated instances, Qwen3-4B, Qwen3-8B and Olmo-3-7B-Instruct achieve an average of 22.9% improvements across all 18 in-distribution query types, and 8.1% over state-of-the-art causal data generation frameworks on 7 established causal
benchmarks outside the training distribution. More importantly, in real-world medical understanding, legal decision, and tabular
reasoning, UniCo-trained models consistently display more faithful
reasoning traces, outperforming the base models by an average of 20.2% in faithfulness metrics. These suggest that causality-centered training not only strengthens causal
reasoning, but also equips LLMs with a causal mindset in general
reasoning tasks.
Read the full abstract ▾
All authors Qirun Dai, Xiao Liu, Jiawei Zhang, Dylan Zhang, Hao Peng, Chenhao Tan
Machine Learning
Cluster Frequency Conformal Prediction for Local Coverage
Tomer Lavi, Bracha Shapira and Nadav Rappoport
Conformal prediction provides distribution-free coverage guarantees, but in many-class classification it may still under-cover specific classes or subpopulations, preventing safe deployment in high-stakes applications. We propose Cluster Frequency Conformal Prediction (CFCP), a plug-in framework that adapts conformal prediction to local structure in a learned representation space. CFCP clusters learned
embeddings, estimates cluster-level label-frequency distributions from calibration data, and for each test point constructs a sample-specific probability vector by softly mixing nearby cluster distributions regularized with global-prior and reliability-aware shrinkage. This vector is then conformalized using standard set constructors. In the disjoint-split regime, CFCP inherits standard finite-sample marginal validity. Under additional assumptions, CFCP further admits a local-validity interpretation. Since representation clusters aggregate locally similar samples, their empirical class frequencies provide a stable estimate of local label ambiguity. Across image and text
benchmarks, CFCP achieves the best class coverage in 15/16
dataset/score-family comparisons and a competitive prediction set size efficiency, with several settings substantially more efficient. Overall, our results show that cluster-frequency information provides an effective localized signal for improving classwise reliability in many-class conformal prediction.
Read the full abstract ▾
Machine Learning
A comparative study of accuracy and rollout stability of temporal surrogate models
Rajarshi Biswas
Temporal surrogate models are effective for predicting chaotic dynamical systems where computational cost can be prohibitive. Several deep neural network architectures can be used for such purposes. In this work, a few commonly used architectures are compared using a common training protocol. The objective is to fairly assess the impact of model architectures for long-horizon prediction stability. Experiments are carried out for three problems, the double pendulum, the Kuramoto-Sivashinsky equations, and the Kolmogorov flow. The experiments are carried out with matching model capacity. Analysis is also carried out for a scenario where each model is individually optimized. It is observed that in both scenarios, the models exhibit categorical differences in long-horizon rollouts. For a concrete quantification, stepwise error injections and perturbation amplifications are analyzed using metrics such as local jacobian, relative one-step bias, and finite-time Lyapunov growth. Additionally, an attractor analysis is also conducted to assess how well the learned models replicate the underlying system geometry. An ablation study to isolate the impact of each component of a continuous-update architecture is also carried out. It is concluded that models that having integrator-like updates show lower bias and perturbation amplification yielding stable long-horizon rollout and more accurate predictions.
Read the full abstract ▾
Machine Learning
Unifying Value Alignment and Assignment in Cross-Domain Offline Reinforcement Learning with Heterogeneous Datasets
Zhongjian Qiao, Jiafei Lyu, Chenjia Bai et al. (+3 more)
Cross-domain offline reinforcement learning (RL) aims to learn a policy in the target domain with a limited target domain
dataset and a source domain
dataset that exhibits a dynamics shift. Training directly on the original source
dataset typically leads to performance collapse. Recent studies perform data filtering from the perspective of dynamics
alignment or value
alignment to enable efficient policy transfer. However, these studies are typically validated on single-domain or single-behavior-policy source
datasets. In this work, we explore a more general heterogeneous cross-domain offline RL setting, where the source
datasets may be collected from multiple source domains by diverse behavior policies. We first uncover a critical yet overlooked issue in this setting: value misassignment. Empirically and theoretically, we demonstrate that value misassignment can undermine value
alignment, mislead data filtering toward selecting suboptimal samples, and loosen the suboptimality gap, thereby degrading the
agent's performance. To address this issue, we propose V2A, which integrates dynamics
alignment, value
alignment, and value assignment. V2A first employs temporally-consistent modality representation learning to extract dynamics modalities from the source
dataset, followed by modality-aware advantage learning to rectify value
alignment. Finally, it adopts a data filtering paradigm to selectively share source data for policy learning. Empirical results show that V2A significantly outperforms strong baseline methods under general heterogeneous cross-domain offline RL settings.
Read the full abstract ▾
All authors Zhongjian Qiao, Jiafei Lyu, Chenjia Bai, Peisong Wang, Siyang Gao, Shuang Qiu
Machine Learning · Cross-listed from eess.SY
DBPnet: Damper Characteristics-Based Bayesian Physics-Informed Neural Network for Wheel Load Estimation
Tianyi Wang, Tianyi Zeng, Zimo Zeng et al. (+8 more)
Advanced driver assistance systems (ADAS) play an important role in modern automotive intelligence, significantly enhancing vehicle safety and stability. The performance of ADAS critically relies on accurate and reliable vehicle state estimation, particularly from vehicle dynamic sensors. Among these signals, wheel load is a key variable for chassis control and safety-critical functions, yet it remains difficult to estimate robustly due to complex suspension geometry, nonlinear dynamics, and measurement noise. To address this issue, we propose DBPnet, a
Bayesian physics-informed neural network (PINN) with a physics-aware
embedding module inspired by damper characteristics. First, this paper presents a suspension linkage-level modeling (SLLM) approach that constructs a nonlinear instantaneous dynamic model by explicitly considering the complex
geometric structure of the suspension. Building upon SLLM,
Bayesian inference is integrated into the PINN to effectively cope with noise and uncertainty in the vehicle chassis system, thereby improving the model's robustness. Then, a physics-informed loss function is employed to ensure consistency with fundamental physical principles, while the damper characteristics-inspired
embedding module extracts temporal variation features of input signals and incorporates them into each layer of the PINN, ensuring that physical observations guide the neural network without being constrained by fixed physical models. Extensive evaluations on high-fidelity simulations and real-world experiments demonstrate that our DBPnet consistently achieves lower RMSE and MaxError than baseline methods. These results highlight the potential of our DBPnet to advance wheel load estimation and contribute to the development of more reliable ADAS actuator functions.
Read the full abstract ▾
All authors Tianyi Wang, Tianyi Zeng, Zimo Zeng, Feiyang Zhang, Yujin Wang, Xiangyu Li, Yiming Xu, Sikai Chen, Junfeng Jiao, Christian Claudel, Xinbo Chen
Machine Learning
The Concept Allocation Zone: Tracking How Concepts Form Across Transformer Depth
James Henry
Concept formation in
transformer language models is a depth-extended process, not a single-layer event: a concept becomes separable across one or more contiguous regions of the residual stream - its Concept Allocation Zone (CAZ). A CAZ is not a concept but the depth segment where the model organizes its geometry to make one separable - concepts may share a CAZ, and typically span multiple across depth; the companion GEM paper shows the separating direction continues to rotate within a CAZ before stabilizing past its boundary. We formalize the CAZ through three layer-wise metrics - Separation, Concept Coherence, and Concept Velocity - with automated boundary detection that applies no significance threshold to CAZ membership (every segment is a CAZ; "strong" vs. "gentle" is score, never a binary cut). Empirical validation across 35 models, 8 architectural families, and 7 concepts shows the separation curve S(l) is frequently
multimodal, and scored detection surfaces a further category of subtle allocation regions ("gentle CAZes") invisible to standard peak detection. The framework generates seven testable predictions; its contribution is the instrument and the phenomena it surfaces - the scored detector, the three metrics, and the
multimodal/gentle-CAZ findings - not the predictions themselves. Released as the open-source rosetta_tools library (v1.3.1).
Read the full abstract ▾
Machine Learning
T2S-MPC: Time-Embedded Online Adaptive Model Predictive Control for Time-Varying Dynamics
Zeyu Shen, Zhuoyuan Wang and Laixi Shi
Recent advances in learning-based model predictive control (MPC) have leveraged neural networks for online model learning, achieving strong performance when nonstationary system dynamics deviate from nominal models. However, existing approaches primarily address specific or relatively structured forms of dynamical variation, leaving more general, unknown, and unpredictable time-varying dynamics insufficiently handled. To tackle this challenge, we propose T2S-MPC, a framework that adaptively learns a residual dynamics model online and integrates it with the nominal model within the MPC framework to enable fast-evolving online planning. To make the model time-aware, we explicitly encode temporal information through a structured time
embedding and employ a two-timescale update scheme, allowing the controller to capture nonstationary dynamics while balancing rapid adaptation with stable learning. We evaluate the proposed method on a 2D quadrotor across stabilization and trajectory tracking tasks under diverse time-varying disturbances, including linear drifting and periodic perturbations. Experimental results show that T2S-MPC consistently outperforms classical MPC, neural MPC, and ablated variants in control performance, while also demonstrating strong robustness across a wide range of disturbance conditions without additional tuning. The source code is publicly available at https://github.com/Zeyuu0920/T2S_MPC
Read the full abstract ▾
Machine Learning
Tiny Brains, Giant Impact: Uncovering the Keystone Neurons of LLM with Just a Few Prompts
Xiangtian Ji, Yuxin Chen, Zhengzhou Cai et al. (+3 more)
Large
language models (LLMs) display strong comprehensive abilities, yet the internal mechanisms that support these behaviors remain insufficiently understood. In this work, we show that across a wide range of open-weight
Transformers, a subset of
neurons remains consistently highly activated during inference across tasks of multiple capability dimensions. By probing along the cross-task activation strength, an extremely sparse subset is isolated, whose removal causes a collapse in model behavior, which we term keystone
neurons. Our analysis reveals that keystone
neurons are a stable and intrinsic
neuron subset of the model that is largely established during pretraining. The parameters associated with these
neurons are tightly calibrated during the training process, and their precise values are critical for the capabilities of the model. Building on these insights, we propose a supervised fine-tuning approach that updates only keystone
neurons, achieving task gains comparable to or even better than full-parameter fine-tuning while better preserving performance in other capability dimensions, despite modifying a much smaller number of parameters.
Read the full abstract ▾
All authors Xiangtian Ji, Yuxin Chen, Zhengzhou Cai, Xiang Wang, An Zhang, Tat-Seng Chua
Machine Learning
DriftingMol: Decoder-Coupled Drift for One-Pass Property-Conditional Molecular Generation
Jiangjie Qiu, Yijun Li, Wentao Li et al. (+1 more)
Property-conditional molecular generation should produce valid, diverse molecules while responding to continuous target values at low sampling cost. We introduce DriftingMol, a two-stage framework that adapts drifting models to a SELFIES latent molecular space. A frozen SELFIES beta-VAE provides the latent space, and the hidden representation of its decoder serves as the drift feature map. In decoder-coupled drift, decoder weights remain fixed, but drift gradients are backpropagated through the decoder feature map to a DiT generator, inducing a pullback metric aligned with molecular decoding. On ZINC250K, the default setting achieves QED Spearman correlation 0.493 with 94.7% uniqueness, while the strongest decoder-coupled condition reaches 0.510. Under protocol-matched four-property conditioning, decoder-coupled drift reaches mean Spearman correlation up to 0.598. Across 15 controlled variants, models that preserve the gradient path through decoder features achieve higher correlations than the tested latent-space, random-feature, and external-feature drift variants, while detached or stop-gradient decoder controls yield near-zero QED correlation and very low uniqueness. These results indicate that decoder-coupled drift is a useful low-cost mechanism for property-biased molecular generation, requiring one generator evaluation and one frozen decoder pass.
Read the full abstract ▾
All authors Jiangjie Qiu, Yijun Li, Wentao Li, Xiaonan Wang
Machine Learning · Cross-listed from stat.ME
Correcting test set contamination by spiking the training data
Johnny Tian-Zheng Wei, Jerry Li, Ameya Godbole et al. (+1 more)
The literature on test set contamination largely focuses on detection, but the correction of contaminated test scores is underexplored. Our core proposal is to spike the training data by intentionally contaminating some test examples at known rates. The spiked examples can then be used to calibrate predictors of model memorization which enable principled statistical correction of inflated test scores. To evaluate different correction
estimators, we first present a simulation framework based on the Hubble models. Hubble models come in minimal pairs, where the perturbed model was deliberately contaminated with several test sets, while the standard model was not, serving as the counterfactual and correction target. We consider
estimators that use information from a memorization predictor, correctness predictor, or both. In simulation, we establish basic statistical intuitions and show that
estimators leveraging memorization and correctness information are better than naive estimation which makes no correction at all. We then instantiate several memorization and correctness predictors, and find that simple predictors such as Platt-scaled membership inference metrics provide good signal for correction. Finally, we examine the practical considerations of spiking. Simple memorization predictors need no more than 10 examples for calibration and often transfer from one
dataset to another. Taken together, spiking is a promising solution for test set contamination.
Read the full abstract ▾
All authors Johnny Tian-Zheng Wei, Jerry Li, Ameya Godbole, Robin Jia
Machine Learning · Cross-listed from cs.CR
RouteScan: A Non-Intrusive Approach to Auditing MoE LLMs Safety via Expert Routing Telemetry
Bo Lv, Zhiheng Xu, KeDong Xiu et al. (+4 more)
As Mixture-of-Experts (MoE) architectures are increasingly adopted for scaling Large
Language Models (LLMs), safety
auditing becomes necessary to verify whether these models produce or facilitate harmful behaviors during operation. However, existing content-based
auditing methods typically require access to user
prompts, model internals, or outputs, potentially exposing sensitive user information and creating a tension between LLM safety and user
privacy. On the other hand, we observe that, in MoE models, different inputs induce different sparse expert-routing patterns, which produce measurable footprints in low-level GPU execution telemetry. We refer to these hardware-observable signals induced by expert-routing decisions as expert routing telemetry; they are derived from GPU execution rather than from router logits or
token-level routing assignments. Inspired by this observation, we propose RouteScan, a non-intrusive
auditing framework for detecting harmful behaviors through such routing-induced GPU telemetry. Specifically, RouteScan utilizes the number of active GPU threads allocated to expert modules during the prefilling phase as a discriminative micro-architectural fingerprint, and builds a lightweight detection pipeline that isolates cross-domain invariant risk indicators for the precise identification of malicious
prompts. Comprehensive evaluations on four open-source MoE LLMs with distinct routing designs demonstrate that RouteScan achieves strong generalization, with an AUROC exceeding 0.91 on unseen harmful domains. Moreover,
privacy stress tests show that, although aggregated execution telemetry retains input-related attribute information, full
prompts and exact sensitive fields cannot be reliably recovered under the evaluated attacks.
Read the full abstract ▾
All authors Bo Lv, Zhiheng Xu, KeDong Xiu, Ruyi Ding, Tianhang Zheng, Zhibo Wang, Kui Ren
Machine Learning
Cross-Domain Energy-Guided Diffusion Generation for Off-Dynamics Reinforcement Learning
Yu Yang, Yihong Guo, Anqi Liu et al. (+1 more)
Off-dynamics offline reinforcement learning seeks to learn a target-domain policy from a large source
dataset and a limited target
dataset under mismatched transition dynamics. Existing approaches such as reward augmentation and data filtering are constrained to the source
dataset and cannot synthesize new target behavior to improve coverage beyond the collected source trajectories. While recent model-based methods attempt to address this by learning target-aware dynamics, the generated experience is constructed only at the transition level, which leads to accumulated errors over long horizons. These limitations necessitate a shift toward trajectory-level generation for off-dynamics offline RL. We propose CEDGE, a Cross-domain Energy-guided
Diffusion GEneration framework. CEDGE trains a trajectory
diffusion model on source-domain trajectories and adapts the generated samples to the target domain through energy guidance. This guidance is derived by minimizing the distribution mismatch between the source and desired target-domain trajectories and is decomposed into return, domain, and behavior energy components. The resulting energy-guided trajectories are useful both for direct planning and as synthetic data for policy learning. Since target adaptation is achieved via energy guidance rather than retraining the
diffusion model, CEDGE can be efficiently adapted to new target dynamics compared to previous methods. Experiments on the ODRL
benchmark demonstrate that trajectory-level energy-guided generation improves
diffusion planning under dynamics shifts and produces synthetic data that improves downstream target policy learning.
Read the full abstract ▾
All authors Yu Yang, Yihong Guo, Anqi Liu, Pan Xu
Machine Learning
Disentangled Double Machine Learning for Accurate Causal Effect Estimation
Guodu Xiang, Kui Yu, Yujie Wang et al. (+3 more)
Confounding bias is a key challenge in causal effect estimation from observational data. Double Machine Learning (DML) addresses this issue by estimating treatment and outcome nuisance functions, constructing treatment and outcome residuals, and estimating causal effects from the residuals. However, DML often produces biased and unstable estimates in highdimensional or finite-sample scenarios. One reason is that DML estimates nuisance functions using all
covariates without disentangling distinct latent factors, resulting in unreliable nuisance function estimation. Another is that imprecise nuisance estimation further introduces residual dependence between the treatment residual and the remaining outcome error, undermining the accuracy of causal effect estimates. To address these issues, in this paper, we propose Disentangled Double Machine Learning (DDML), a novel
algorithm that integrates two key strategies. First, a causal role disentanglement strategy decomposes
covariates into confounders, treatment-specific factors, and outcomespecific factors for enabling reliable nuisance function estimation. And second, a residual dependence orthogonalization strategy mitigates residual dependence caused by nuisance estimation errors for enhancing the precision of causal effect estimates. Experimental results on synthetic, semi-synthetic, and real-world
datasets demonstrate that DDML significantly outperforms 13 state-of-the-art baseline
algorithms in both MAE and RMSE.
Read the full abstract ▾
All authors Guodu Xiang, Kui Yu, Yujie Wang, Richang Hong, Fuyuan Cao, Jiye Liang
Machine Learning
Active Learning for Stochastic Contextual Linear Bandits
Emma Brunskill, Ishani Karmarkar and Zhaoqi Li
A key goal in
stochastic contextual linear bandits is to efficiently learn a near-optimal policy. Prior
algorithms for this problem learn a policy by strategically sampling actions but naively (passively) sampling contexts from the underlying context distribution. However, in many practical scenarios – including online content recommendation, survey research, and clinical trials – practitioners can actively sample or recruit contexts based on prior knowledge of the context distribution. Despite this potential for active learning, the role of strategic context sampling in
stochastic contextual linear bandits is underexplored. We propose an
algorithm that learns a near-optimal policy by strategically sampling rewards of context-action pairs. We prove instance-dependent theoretical guarantees demonstrating that our active context sampling strategy can improve over the minimax rate by up to a factor of $\sqrt{d}$, where $d$ is the linear dimension. We show empirically that our
algorithm reduces the number of samples needed to learn a near-optimal policy, in tasks such as warfarin dose prediction and joke recommendation.
Read the full abstract ▾
Machine Learning · Cross-listed from math.OC
Lifted Schrödinger Bridges for Gaussian Mixture Endpoints: Projection Gaps and Path-Space Obstructions
Siddhartha Ganguly, George Rapakoulias and Panagiotis Tsiotras
We study
stochastic density control between Gaussian-mixture endpoint distributions under Brownian prior dynamics. Since the direct Schrödinger bridge between Gaussian mixtures is generally not available in closed form, we introduce a lifted path-space construction in which each trajectory is augmented with a source–target component label. Consequently, the problem decomposes into Gaussian component-to-component Schrödinger bridges with explicit marginal, drift, and cost formulas, while the mixture-level assignment reduces to a finite-dimensional entropic coupling problem with a Sinkhorn scaling form. We then analyze the projection obtained by discarding or forgetting the label. By construction, the projected law satisfies the original Gaussian-mixture endpoint constraints, but its relative
entropy generally differs from the lifted relative
entropy by a nonnegative conditional label-information gap. This gap reveals a path-space obstruction: the lifted optimizer cannot, in general, be identified with the direct unlabeled Schrödinger bridge after projection. We also derive the posterior-averaged Markov drift associated with the projected marginal flow, prove a kinetic-energy upper bound, and identify a common path-potential condition under which the projection gap vanishes. Several numerical illustrations showing density and shape control are recorded for a self-contained exposition.
Read the full abstract ▾
Machine Learning
CONF-KV: Confidence-Aware KV Cache Eviction with Mixed-Precision Storage for Long-Horizon LLM
Yubo Li and Yidi Miao
Long-horizon LLM inference turns the key–value (KV) cache into the dominant GPU
memory consumer and makes per-
token attention increasingly expensive. Many common eviction policies use static recency windows or historical
attention, leaving unused a signal computed on every decoding step: the model's current uncertainty. We introduce CONF-KV, a KV-cache manager that converts the next-
token distribution into a scalar confidence score and uses it to choose the per-step cache budget, retaining more context when the model is uncertain and
pruning aggressively when it is confident. Within each budget,
tokens are ranked by a composite of accumulated
attention mass and recency, while a protected recent window preserves local coherence. We combine the policy with blockwise online-softmax
attention, mixed FP16/INT8 storage, and a pyramidal per-layer budget variant. Across four model families and generated lengths up to 4K, CONF-KV stays near the footprint of a fixed 512-
token sliding window while remaining within 1.5–2.1 perplexity points of full KV. On Needle-in-a-Haystack up to 32K
tokens, CONF-KV reaches 91.4%
retrieval accuracy versus 53.8% for sliding windows and 80.6% for H2O; on 75 VisualWebArena tasks it retains 95.3% of full-KV success at 2.8 times lower peak
memory.
Read the full abstract ▾