Machine Learning
RECAST: Model Reconstruction via Counterfactual-Aware Wasserstein Geometry under Limited Data
Xuan Zhao, Lena Krieger, Zhuo Cao et al. (+3 more)
Counterfactual explanations (CFs) help understand machine learning models by identifying minimal input changes that would lead to alternative model outcomes. Recent work demonstrates their utility for reconstructing black-box models, enabling third-party
auditing of opaque decision systems for
fairness and accountability. Still, CF-based reconstruction may suffer from decision boundary shifts, overfitting, and restrictive assumptions requiring online query access to target platforms. We propose REconstruction via Counterfactual-Aware waSserstein opTimization (RECAST) under limited data and restricted access, a behavioral surrogate model based on Wasserstein barycentric prototypes. Our approach addresses decision boundary shifts by incorporating CFs as informative, though less representative, samples for both classes, maintaining high surrogate fidelity in low-sample regimes without requiring online access during reconstruction. To enhance
fairness auditing, our method enables systematic group
fairness diagnostics. Experiments on real-world
datasets and various setups show that RECAST effectively achieves high fidelity and query efficiency, as well as stable results even when the access is limited and noisy.
Read the full abstract ▾
All authors Xuan Zhao, Lena Krieger, Zhuo Cao, Arya Bangun, Hanno Scharr, Ira Assent
Machine Learning · Cross-listed from cs.CL
VASAE: Naming SAE Dictionary Directions with Vocabulary-Aligned Anchoring
Kairui Zhang, Ziwen Yu, Zahraa S. Abdallah et al. (+1 more)
Sparse autoencoders (SAEs) provide useful decompositions of
Transformer residual streams, but their learned features are usually named post hoc rather than directly connected to the
Transformer's
token vocabulary. We introduce Vocabulary-Aligned Sparse Autoencoder (VASAE), a method that trains SAE features under vocabulary-aligned anchoring and assigns each feature an intrinsic
token name: the
token string whose
embedding is nearest to that feature. Without reducing reconstruction quality compared with a standard SAE, VASAE produces dictionaries with vocabulary-aligned features. Using a 0.8 cutoff on the nearest-
token alignment score, dictionaries trained on GPT-2-small post-residual streams align about 90% of features in layers 0–10. In Llama-3.1-8B, representative shallow and middle-layer dictionaries contain strongly aligned features, including 92.8% in the shallow layer, while the representative final-layer dictionary shows limited
alignment. After subtracting the sentence-level mean sparse code, case studies show that many remaining intrinsic
token names are relevant to nearby input
tokens. These results suggest that vocabulary-aligned anchoring can connect learned features to intrinsic
token names during training, complementing post hoc interpretation of learned dictionaries.
Read the full abstract ▾
All authors Kairui Zhang, Ziwen Yu, Zahraa S. Abdallah, Martha Lewis
Machine Learning
Two-Stage Fine-Tuning for Protein Sequence Generation with Targeted Amino-Acid Composition
Violeta Basten-Romero, Rubén Muñoz-Tafalla, Anna María Díaz-Rovira et al. (+3 more)
Protein language models are standard priors for biological sequence generation, but steering them toward explicit distributional design targets remains largely unexplored. We study a constrained
protein generation problem in which sequences must match a desired amino-acid (AA) composition profile while preserving plausible sequence statistics and diversity. The motivating application is synthetic feed
protein design, where the AA composition of dietary
proteins directly determines their nutritional value. We propose a two-stage pipeline in which domain-adaptive fine-tuning (FT) on an in-domain
protein dataset is followed by iterative reward-weighted FT via reinforcement learning (RL) anchored against the FT model as a frozen reference. We evaluate the pipeline on two AA compositions and find that FT brings the average composition close to the target, while the subsequent RL enforces specific sequence constraints that FT alone cannot satisfy. We additionally evaluate the design choices of the proposed composition reward term against two baselines and an ablated variant, isolate the contribution of each training stage, and verify that AA composition
alignment is achieved without degrading sequence quality.
Read the full abstract ▾
All authors Violeta Basten-Romero, Rubén Muñoz-Tafalla, Anna María Díaz-Rovira, Bertran Miquel-Oliver, Isaac Filella-Merce, Víctor Guallar
Machine Learning
Graph Dimensionality Reduction for Contextual Bandits: Structure-Specific Regret Bounds under Approximate Smoothness and Noisy Eigenspaces
Joyanta Jyoti Mondal, Ibne Farabi Shihab and Anuj Sharma
Contextual bandits with graph-structured arms arise in recommendation, citation
retrieval, and social advertising, where arms connected on a graph tend to share reward signal. Standard dimensionality reduction ignores this structure, inflating exploration cost by a factor of $d/k$. We propose GraphDR-LinUCB, which projects arm features onto the graph's low-frequency spectral subspace and runs linear UCB in the resulting $k$-dimensional space. We prove the first $\wtO(k\sqrt{T})$ regret bound for spectral-projection-based contextual bandits, reducing dimension dependence from $d$ to $k$; a perturbation argument extends this to noisy graphs, with an explicit penalty for reward-smoothness mismatch and graph-estimation error. Our central theoretical finding is that the high-frequency reward component need not incur a worst-case linear-in-$T$ penalty: its actual cost depends on its realized impact along the played path, not on its total energy. A simple spectral comparison between subspaces ($\Gamma_k$) predicts which reducer wins on a given
dataset, correctly calling five of six real-
dataset outcomes without any fitted threshold. Across a synthetic
benchmark and six real
datasets (MovieLens, Amazon, LastFM, ogbn-arxiv, MIND), GraphDR-LinUCB reduces cumulative regret by $15\times$ over full-dimensional LinUCB and outperforms competing graph-aware methods on five of six; the single failure is precisely where the graph's spectral subspace is misaligned with the reward.
Read the full abstract ▾
Machine Learning
TA-SparseMG: Trend-Aware Sparse Forecasting via Multi-Scale Gating for Long-Term Time Series
Wenchao Liu, Hongbing Wang, Youji Zhu et al. (+2 more)
Long-term time series
forecasting finds extensive applications in domains such as power demand, traffic flow, meteorological observation, and renewable energy dispatch.
Forecasting dynamically varying long-term time series poses inherent challenges, including statistical nonstationarity, local high-frequency disturbances, and coupled cross-period dependencies, which make it difficult for lightweight models to balance parameter efficiency and
forecasting performance. To address this issue, this study presents TA-SparseMG, a lightweight cross-period
forecasting model built on SparseTSF's sparse cross-period modeling framework. It incorporates three key modules: a trend-aware reversible instance normalization module, a scale-adaptive gated denoising module, and a multiscale gated-
attention MLP
forecasting module. The trend-aware normalization module captures input-window statistics and calibrates forecast-window distributions, effectively mitigating distribution shift. The scale-adaptive gated denoising module performs feature smoothing and residual suppression before period rearrangement, thereby reducing interference from high-frequency perturbations. The multiscale gated
attention prediction module strengthens the prediction head's adaptive representational capacity via conditional gating and feature modulation. Extensive experiments across multiple LTSF
benchmarks demonstrate that the proposed TA-SparseMG consistently achieves superior, stable performance. Ablation studies confirm that each module independently improves distribution adaptation, input robustness, and cross-period feature mapping capability.
Read the full abstract ▾
All authors Wenchao Liu, Hongbing Wang, Youji Zhu, Xiaodong Liu, Xiangguang Xiong
Machine Learning · Cross-listed from physics.comp-ph
Mosaic: A Benchmark Suite for Differentiable Physics Solvers
Andrin Rehmann, Heiko Zimmermann and Dion Häfner
Differentiable partial differential equation (PDE) solvers underpin solver-in-the-loop ML training, gradient-based optimal control, and inverse problems, yet the practical cost of obtaining correct, usable gradients from a given solver on a given problem is largely undocumented. Integration effort, computational cost, gradient accuracy, and numerical conditioning vary widely across solvers and are discoverable only by trial and error. We introduce Mosaic, an extensible benchmarking framework for differentiable PDE solvers that standardizes access to solver gradients. Each solver is packaged as a containerized component (Tesseract) exposing a uniform gradient API regardless of
language or automatic differentiation (AD) strategy, enabling researchers to evaluate, compare, and build on non-trivial physical solvers. Our evaluation of 14 solvers across fluid dynamics, structural mechanics, and heat transfer demonstrates that the
benchmark surfaces practically relevant differences: order-of-magnitude variation in computational cost and Jacobian conditioning, alongside structural incompatibilities that eliminate solvers from realistic tasks entirely. Despite this variation, all solvers that produce gradients converge to similar optima, indicating that the practical barriers are
memory limits, numerical stability, and setup compatibility rather than gradient accuracy alone. Mosaic is open-source and available at https://github.com/pasteurlabs/mosaic.
Read the full abstract ▾
Machine Learning
A Comparison of Fusion Techniques for Multi-Modal Human Activity Recognition on the HARMES Dataset
Ahmed Mohamady, Robin Burchard and Kristof Van Laerhoven
Recent advances in Human Activity Recognition (HAR) from wearable sensors have shown that multi-modal deep learning models consistently outperform their uni-modal counterparts. Modalities can include IMUs, RGB cameras, audio signals, and others. One important aspect of multi-modal deep learning is the sensor fusion approach we apply. Over recent years, multiple fusion paradigms have been proposed for multi-modal HAR. However, to the best of our knowledge, no head-to-head comparison of these paradigms exists on a common multi-modal HAR
benchmark dataset. To address this research gap, we systematically compare seven state-of-the-art sensor fusion methods on the recently released HARMES
dataset, which comprises 61 hours of fully labeled IMU, audio, and ambient humidity data. The chosen
dataset focuses on 15 household and personal hygiene activities of daily living (ADLs). By applying the seven different fusion techniques to a state-of-the-art multi-modal model architecture, we show that Gated Multi-modal Fusion achieves the highest macro F1-score (0.82), surpassing the concatenation-based late fusion HARMES paper baseline of 0.76 by +6pp under leave-one-participant-out evaluation. All code used in our experiments is made publicly available on GitHub.
Read the full abstract ▾
Machine Learning
FlexMoE: One-for-All Nested Intra-Expert Pruning for MoE Language Models
Fan Mo, Yuxuan Han, Geng Zhang et al. (+2 more)
Mixture-of-Experts (MoE)
language models scale model ability with sparsely activated experts, making this architecture a standard recipe for modern large models. However, sparse activation does not remove the deployment burden of storing and serving all experts, and the available deployment budget can vary substantially across devices, users, and workloads. Existing MoE compression methods are still largely fixed-budget, typically optimizing one compressed endpoint at each chosen target budget. We study a different setting: converting a large pretrained MoE LLM into a nested family of deployable subnetworks across budgets. Our method first ranks expert FFN channels by their importance, then lets each expert learn a discrete action to prune its channels. By gradually increasing cost pressure, a single action-training run exports a series of action masks from high to low budgets, each of which identifies a reliable smaller subnetwork nested in the ranked base model. Moreover, we use a single recovery fine-tune at a mid
pruning budget (40%) to recover degraded model quality and transfer the recovered model to other unseen budgets. Overall, our framework surpasses recent MoE compression baselines. Specifically, on Qwen2-57B-A14B, our method retains ~99.8% of base performance while
pruning 50% of routed expert parameters even without fine-tuning. For deployment, our pruned subnetworks deliver real
memory reduction and throughput gains, and further support realtime online budget switching with kernel-level co-design.
Read the full abstract ▾
All authors Fan Mo, Yuxuan Han, Geng Zhang, Wangbo Zhao, Yang You
Machine Learning · Cross-listed from cs.CV
A Unified Framework for Vision Transformers Equivariant to Discrete Subgroups of $\mathrm{O}(2)$
Tīkun Ông and Georg Bökman
Vision transformers have become a dominant architecture for visual recognition. However, standard models do not explicitly encode the planar
symmetries that arise in many
vision domains. We introduce a family of
vision transformers equivariant to arbitrary discrete subgroups of $\mathrm{O}(2)$, providing a unified framework that generalizes prior flipping- and $D_4$-equivariant
transformer architectures. Our construction yields equivariant analogues of the core
transformer components, together with expressivity guarantees for the resulting layers. In particular, we show that whenever $H \le G$, the class of $G$-equivariant ViTs embeds naturally into the class of $H$-equivariant ViTs. We also prove that, in the single-head setting, the corresponding equivariant self-
attention layer realizes every $G$-equivariant self-
attention map representable by ordinary self-
attention. We further construct a $D_6$-equivariant model based on hexagonal patches, making the architecture compatible with six-fold rotational
symmetries. We evaluate the resulting models on the PatternNet aerial image
dataset in artificially data-scarce regimes across subgroups of $D_4$ and $D_6$. Our experiments compare two equivariant
attention mechanisms and analyze how the choice of homogeneous-space configurations used in the nonlinearities affects performance. Preliminary results under matched parameter budgets indicate that equivariance can improve recognition accuracy, motivating further study of how discrete
symmetry groups shape
transformer-based visual recognition models.
Read the full abstract ▾
Machine Learning
GNBAN: Graph Neural Basis Attention Networks for Long-Horizon Forecasting over Large Entity Sets
Janak M. Patel, Anirudh Deodhar and Dagnachew Birru
Demand
forecasting at the bottom of a retail hierarchy requires predicting tens of thousands of correlated long-horizon series across products, stores, and regions. Modern systems must scale across massive catalogs, capture shared demand dynamics, and remain interpretable enough to be trusted. Classical statistical methods need a separate model per series and are hard to manage at scale; deep autoregressive models struggle as the joint state grows to tens of thousands of dimensions; and recent graph-based forecasters, while capturing cross-entity dependencies, often produce opaque long-horizon forecasts. We propose GNBAN (Graph Neural Basis
Attention Network), an end-to-end architecture combining heterogeneous graph representation learning with an interpretable basis-decomposition head. Retail data are represented directly as a heterogeneous graph derived from the relational schema, so a single model serves the entire catalog. Rather than predicting the horizon directly, GNBAN decomposes each forecast into trend, seasonal, and generic components. Its key innovation is a per-basis
attention mechanism: each basis function keeps its own learnable query and retrieves information independently from the entity's historical neighborhood, letting different bases specialize to distinct temporal patterns while preserving interpretability. On two large-scale
benchmarks, M5 Walmart and Favorita Grocery Sales, evaluated under matched protocols, GNBAN improves volume-weighted WRMSSE by roughly 4-5% over a matched graph baseline. Qualitative analysis shows the learned decomposition exposes trend, seasonal, and residual demand drivers without post-hoc explanation methods. These results demonstrate that scalable relational
forecasting and interpretable forecast decomposition can be achieved together in a unified graph-based framework.
Read the full abstract ▾
Machine Learning
Applicability of memorization indicators for early spotting of overfitting while recalibrating sEMG-decoders on low sample sizes
Stephan J. Lehmler, Tobias Glasmachers and Ioannis Iossifidis
Deep learning models for surface electromyography (sEMG) can benefit substantially from subject-specific (re-)calibration, since no sufficiently large and diverse
datasets are available to train fully generic decoders. However, for user acceptance, the number of repetitions that can realistically be collected during calibration is severely limited, which increases the risk of overfitting and, in extreme cases, can even degrade performance compared to the uncalibrated model. Classical overfitting indicators such as validation performance and regularization with early stopping are difficult to apply in this low-sample regime, as they require additional held-out data that is rarely available in practical calibration scenarios. In this work, we investigate a recently proposed class of memorization indicators based solely on the activation statistics of rectified linear units (ReLU) in deep neural networks, which can be computed directly from training data without any extra validation set. We conduct a transferlearning experiment on a
benchmark sEMG
dataset, where a convolutional neural network is first pre-trained on multiple subjects and subsequently fine-tuned on individual users using only a small number of repetitions. During calibration, we monitor both decoding performance and the activation behaviour of the last hidden layer. Our results provide first evidence that decreases in test accuracy during fine-tuning are ac companied by characteristic changes in activation rates, indicating that activation-based memorization indicators are a promising tool for early spotting of unsuccessful learning in low-sample sEMG calibration settings.
Read the full abstract ▾
Machine Learning
WattLayer: Get Layers Right to Estimate Inference Energy of Neural Networks
Adrien Sardi, Marie-Line Alberi Morel, Sara Alouf et al. (+2 more)
The widespread adoption of Artificial Intelligence (AI) has led to increasing concerns about energy consumption, yet there is a lack of standardized methodologies to accurately estimate AI inference energy consumption, particularly across various tasks and architectures. In this study, we propose a task independent, layer-wise energy estimation model for AI architectures. Our model is evaluated on a large
dataset of more than 100,000 layers for 295 neural network architectures across 3 widely-used tasks and 3 distinct hardware platforms. Our approach achieves a median error of 19.6%, outperforming state-of-the-art methods. We further show that layer-wise decomposition generalize to new tasks without complete retraining, by leveraging shared layers across architectures. It offer tools, insights and a precise methodology to empower stakeholders in designing energy-efficient AI systems.
Read the full abstract ▾
All authors Adrien Sardi, Marie-Line Alberi Morel, Sara Alouf, Frédéric Giroire, Joanna Moulierac
Machine Learning
USAD: Uncertainty-aware Statistical Adversarial Detection
Zhijian Zhou, Xunye Tian, Jiacheng Zhang et al. (+5 more)
Statistical
adversarial detection (SAD) treats detection as a two-sample test. Given a reference set of clean examples (CEs) and a batch of queries, potentially containing an unknown mixture of CEs and
adversarial examples (AEs), SAD decides whether the query distribution drifts away from the CE distribution while controlling the false-alarm rate. Existing SAD-based methods mainly use maximum mean discrepancy (MMD) to measure the distributional discrepancy. However, MMD's distributional properties limit its ability to capture characteristic uncertainty patterns of AEs that are crucial for detection: AEs typically exhibit abnormal feature spread (i.e., global uncertainty) and instability under perturbations (i.e., local uncertainty). To close the gap, we propose Uncertainty-aware Statistical
Adversarial Detection (USAD), which explicitly captures these uncertainty patterns with two new statistics: (1) Variance Discrepancy (VD), which measures the difference in feature spread between AEs and CEs to capture global uncertainty differences. (2) Perturbation-based Covariance Discrepancy (PCD), which compares feature covariance under Gaussian perturbations to capture local uncertainty differences. By aggregating VD and PCD, USAD achieves superior detection performances over baseline methods against various
adversarial attacks, highlighting the importance of considering characteristic behaviors of AEs for effective SAD. Our code is available at: https://anonymous.4open.science/r/USAD.
Read the full abstract ▾
All authors Zhijian Zhou, Xunye Tian, Jiacheng Zhang, Zesheng Ye, Yiyi Guo, Donghao Zhang, Liuhua Peng, Feng Liu
Machine Learning
Pepti-drift: Scalable Safe-Active Peptide Generation Without Inference-Time Guidance
Takashi Fujiwara, Hikaru Shindo, Kaushalya Madhawa et al. (+6 more)
Therapeutic peptides are a promising drug modality, but their generation must satisfy multiple therapeutic constraints. We introduce BindSafe-PepBench, a fixed-budget
benchmark that jointly evaluates target binding and four major safety metrics on the same generated candidates. We reveal that peptide length is a major confounder of joint binding-safety evaluation: longer peptides tend toward stronger predicted binding but less favorable predicted safety. This creates an apparent trade-off and can bias comparisons among models with different output-length distributions. We report absolute Safe-Active yield and exact-length-matched gains to distinguish generative improvements from output-length effects. High Safe-Active yield remains challenging, while the strongest multi-property methods rely on costly inference-time guidance. We therefore introduce Pepti-drift, a one-step generation framework that incorporates attraction toward target-specific binders and repulsion from liability-associated regions, requiring a single latent refinement followed by parallel decoding without inference-time guidance. Across 88 held-out targets, Pepti-drift achieves an 18.37% predicted Safe-Active yield while retaining positive exact-length-matched gains. The resulting gains are competitive with multi-property-guided baselines while requiring 468 times lower generation cost, enabling scalable and fair high-throughput peptide design.
Read the full abstract ▾
All authors Takashi Fujiwara, Hikaru Shindo, Kaushalya Madhawa, Jun Jin Choong, Shuan Chen, Yuna Oikawa, Yiming Zhang, Gyubok Lee, Keisuke Ozawa
Machine Learning · Cross-listed from quant-ph
Parameter-Efficient Quantum-Inspired Fast Weight Programmers for Traffic-Matrix Forecasting
Kuo-Chung Peng, Jiun-Cheng Jiang, Chun-Hua Lin et al. (+3 more)
Traffic
matrices (TMs) capture network-wide origin-destination demand and are central to traffic engineering, yet accurate whole-
matrix forecasting remains challenging when prediction must be performed under the
memory, update, and training-budget constraints of online network control. This paper investigates whether compact
quantum-inspired recurrent models can provide effective TM forecasts without relying on dedicated graph,
transformer, or
diffusion modules. We adapt gated
quantum-inspired Kolmogorov-Arnold network fast-weight programmers (QKAN-FWPs) to direct multi-step Abilene TM
forecasting, where each model predicts the next 20 five-minute frames of a 144-channel origin-destination (OD)
matrix from a two-hour history. We
benchmark three QKAN placement variants against a matched-size long short-term
memory (LSTM) network, a larger LSTM, and a classical gated fast-weight programmer under a shared fixed-budget training protocol. Among the evaluated recurrent models, G-QKANFWP achieves the best pooled root-mean-square error (RMSE), while using only 22.4% of the larger LSTM. It also outperforms both the matched-size LSTM and the classical G-FWP baseline, indicating that the gain is not due to gated fast-weight framework alone. Convergence and channel-wise analyses further show that the
quantum-inspired variants obtain lower validation-loss area under the learning curve (AULC) than matched-size recurrent baselines, while G-QKANFWP and GQKAN-FWP achieve substantially more OD-channel wins. These results identify a classical slow programmer with a
quantum-inspired fast programmer as a promising accuracy-efficiency design for resource-conscious network traffic-
matrix forecasting.
Read the full abstract ▾
All authors Kuo-Chung Peng, Jiun-Cheng Jiang, Chun-Hua Lin, Tai-Yue Li, Nan-Yow Chen, Samuel Yen-Chi Chen
Machine Learning · Cross-listed from cs.CV
Scalable and Differentiable Point-Cloud Registration Using Maximum Mean Discrepancy
Rixon Crane, Fahira Afzal Maken, Nicholas Lawrance et al. (+4 more)
We present MMD-Reg, a novel correspondence-free approach to point-cloud registration that is differentiable and has linear computational complexity in the number of points. We model registration as a nonlinear least-squares problem based on the Maximum Mean Discrepancy, approximated using random
Fourier features. The resulting objective can be solved efficiently with standard methods such as Levenberg-Marquardt, and the solution is differentiable via the implicit function
theorem. This allows MMD-Reg to be used as a differentiable optimization layer within end-to-end trainable models, supporting registration under challenging conditions such as poor initial
alignment and partial overlap. We demonstrate this Neural MMD-Reg formulation by integrating the layer with a set
transformer, training the resulting model in supervised and unsupervised settings, and comparing its performance against recent learning-based methods. We also evaluate standalone MMD-Reg, comparing its accuracy and scalability against widely used non-learning-based registration methods.
Read the full abstract ▾
All authors Rixon Crane, Fahira Afzal Maken, Nicholas Lawrance, Stanislav Funiak, Kasra Khosoussi, Ming Xu, Russell Tsuchida
Machine Learning · Cross-listed from quant-ph
Quantum Dynamic Time Warping for Multivariate Time Series Classification
Diego Alvarez-Estevez, Alejandro Mayorga-Redondo and Eduardo Mosqueira-Rey
Dynamic Time Warping (DTW) is a cornerstone for time series classification, but its reliance on Euclidean distances fails to capture latent cross-channel correlations in complex multivariate data. We propose a hybrid
Quantum Dynamic Time Warping (qDTW) architecture, replacing the classical distance metric with the parameterized geometry of a
quantum Hilbert space. Through structural ablation on
benchmarks up to $C=8$ spatial dimensions, we establish fundamental topological rules for
quantum sequence
alignment. We introduce a Unified Pre-
Embedding Adjoint Ansatz that decouples trainable
entanglement from classical data, eliminating the severe phase-scrambling and information bottlenecks inherent to traditional measurements. We demonstrate this decoupled architecture allows untrained
quantum kernels to act as highly expressive baselines, while parameterized training effectively untangles deeply overlapping hyper-dimensional data. Furthermore, we identify a strict spatial-temporal expressivity tradeoff: temporal depth (data re-uploading) is necessary for dimensionally restricted univariate
circuits, but applying it to wide multi-
qubit registers triggers chaotic frequency-spectrum explosions and representation collapse. By navigating these topological hazards, our multivariate
quantum architecture outperforms classical baselines, setting a new standard for integrating parameterized
quantum circuits with dynamic programming
Read the full abstract ▾
Machine Learning
Accelerating Hierarchical Sparse Predictive Coding with Hybrid Amortized Inference
Kazuhisa Fujita
Hierarchical predictive coding provides an interpretable framework for perception as error-driven inference in multi-layer models, while sparse coding imposes parsimonious latent representations through explicit sparsity constraints. Their combination yields hierarchical sparse predictive coding models with appealing computational and neuroscientific properties, but practical use is often limited by the cost of iterative latent inference. In such models, each input may require many recurrent refinement steps before a useful sparse representation is obtained, and this burden becomes more severe as the hierarchy deepens. We study this bottleneck by comparing training-and-inference procedures that share the same hierarchical sparse objective formulation and architecture but use different latent-inference mechanisms. The comparison includes classical iterative inference based on ISTA, an accelerated MFISTA reference, structurally informed amortized inference using a LISTA-style bottom-up encoder adapted to the hierarchical model, and a Hybrid procedure in which this fast amortized initialization is followed by a small number of corrective energy-based refinement steps. Each procedure is trained separately, allowing its inference mechanism to interact with dictionary learning and, where applicable, encoder learning. We measure the resulting reconstruction quality, sparsity, latency, and run-to-run variability across random seeds on static image
benchmarks. The results show that Hybrid improves over pure amortization in the tested settings while remaining substantially faster than procedures based on long iterative inference.
Read the full abstract ▾
Machine Learning · Cross-listed from eess.SP
Distributed Air-Gap Flux and Rotor-Current Fusion for Operating-Regime Identification in a 10-MW Kaplan Hydrogenerator
Eduardo Jr Piedad, Rafel Roig, Xavier Escaler et al. (+2 more)
Reliable monitoring of hydroelectric generators requires descriptors that capture both electrical loading and electromagnetic field behavior. This work investigates operating-regime identification in the Porjus U9 10-MW Kaplan hydrogenerator using synchronized measurements from ten stator-mounted
Hall probes and six rotor-current channels. Seven steady guide-vane-opening settings are considered, and each 300s record is divided into 1s windows. The resulting windows are represented by spatial
Fourier descriptors of the circumferential air-gap field, probe-wise temporal flux indicators, and channel-wise RMS rotor-current features. Correlation analysis and principal component analysis are used to examine how the feature groups vary with the operating point, and Random Forest, radial-basis-function support vector classification, and multilayer perceptron models are evaluated for supervised identification of the guide-vane-opening state. The analysis shows that RMS rotor-current features mainly track the loading axis, while the
magnetic-flux features reveal complementary information associated with spatial imbalance, waveform distortion, and weak low-frequency modulation. Spatial descriptors alone provide limited separability, yielding test accuracies below 27%, whereas rotor-current features alone reach about 84-85%. Combining flux and current information gives the most discriminative representation; the SVC-RBF model achieves 99.5% test accuracy and macro-F1 score. The results indicate that distributed air-gap
magnetic sensing, when fused with rotor-current measurements, can support accurate and interpretable data-driven monitoring of Kaplan hydrogenerator operating regimes.
Read the full abstract ▾
All authors Eduardo Jr Piedad, Rafel Roig, Xavier Escaler, Eduardo Prieto-Araujo, Oriol Gomis-Bellmunt
Machine Learning · Cross-listed from cs.DC
Optimizing Teacher-Student Partitioning for Scalable Knowledge Distillation on HPC Systems
Adrian P. Dieguez, Victor Conchello Vendrell, Alex Batlle et al. (+3 more)
Knowledge
Distillation (KD) enables training smaller student models under the guidance of larger teacher models, and the widely adopted TRL library implements it. Yet, TRL treats both models symmetrically, missing opportunities to exploit their pronounced asymmetry in
memory footprint, and communication requirements. This paper presents an HPC-aware methodology for KD that decouples teacher and student partitioning efficiently. Our approach achieves up to 67% higher samples-per-second than TRL by avoiding unnecessary teacher-model data structures and selecting the best split strategy. We combine vertical and horizontal partitioning of models, deriving an analytical expression that identifies the existence of inflection points between splitting regimes. These results showed that exploiting teacher–student asymmetry through
topology-aware parallelism notably accelerated GKD training on production HPC clusters at our company
Read the full abstract ▾
All authors Adrian P. Dieguez, Victor Conchello Vendrell, Alex Batlle, Vinnam Kim, Jordi Ros-Giralt, Harris Teague
Machine Learning · Cross-listed from cs.CV
Improving Adversarial Robustness via Activation Amplification and Attenuation
Taïga Gonçalves, Yongsong Huang, Tomo Miyazaki et al. (+1 more)
The existence of
adversarial attacks is often attributed to the presence of non-robust features in neural networks. While prior defenses reduce their impact via
pruning, masking, or feature recalibration, we instead propose to jointly learn to amplify and attenuate these signals through a simple activation scaling mechanism. To this end, we introduce Activation Amplification and Attenuation (A3), a lightweight plug-in module that enhances
adversarial robustness with minimal modifications of the activations. A3 dynamically rescales the activations using a learnable mask and a scaling factor derived from the original activation magnitudes. The influence of
adversarial perturbations can be amplified or attenuated using the same learnable parameters by simply flipping the sign of the scaling operation. The amplified signals serve as negative references to construct novel contrastive and ranking loss functions. Experimental analysis shows that learning to degrade the predictions in amplification mode simultaneously improves
adversarial robustness in attenuation mode. Moreover, A3 relies on only a small number of learnable parameters, with most of its behavior being determined by the scaling mechanism rather than additional network capacity. Extensive experiments demonstrate that integrating A3 into different backbones,
datasets, and training methods consistently improves
adversarial robustness while introducing negligible computational and
memory overhead compared to existing plug-in modules. Code is available at: https://github.com/tgoncalv/A3.
Read the full abstract ▾
All authors Taïga Gonçalves, Yongsong Huang, Tomo Miyazaki, Shinichiro Omachi
Machine Learning · Cross-listed from q-bio.NC
CANNs: A Toolkit for Research on Continuous Attractor Neural Networks
Sichao He, Aiersi Tuerhong, Shangjun She et al. (+4 more)
Continuous attractor neural networks (CANNs) are the canonical computational framework for how the
brain encodes continuous variables such as spatial position, head direction, and movement direction, and explain the activity of hippocampal place
cells, entorhinal grid
cells, and head-direction
cells. CANN research, however, is fragmented: most results rest on lab-specific implementations, general-purpose simulators lack CANN-specific abstractions, and the path from spike trains to attractor geometry in real recordings lacks a standardized toolkit. Here, we present a comprehensive open-source toolkit that unifies the full CANN research workflow. It combines three tightly integrated components: 1) canns, a Python library on BrainPy/JAX that provides standardized 1D/2D CANNs, spike-frequency-adaptation variants, grid
cell networks, hierarchical path-integration models, and
brain-inspired attractor architectures, together with curated
datasets, task generators, an analyzer module and trainer modules for biologically plausible plasticity; 2) canns-lib, a Rust acceleration backend delivering hundreds-of-times speedups for spatial-
navigation workloads and modest gains for Ripser-based persistent homology; 3) ASA (Attractor Structure Analyzer), a PySide6 pipeline applying persistent homology and
cohomology to experimental neural recordings to detect ring-like and toroidal attractor signatures in real data. The toolkit ships with full-detail reproducible pipelines that recover recent CANN results including SFA-driven anticipative tracking, theta sweeps in head-direction/place/grid systems, and hierarchical path integration.
Read the full abstract ▾
All authors Sichao He, Aiersi Tuerhong, Shangjun She, Tianhao Chu, Yuling Wu, Junfeng Zuo, Si Wu
Machine Learning
NormGuard: Reward-Preserving Norm Constraints in Flow-Matching Reinforcement Learning
Tianlin Pan, Lianyu Pang, Cheng Da et al. (+4 more)
Reinforcement learning (RL) post-training improves the reward
alignment of flow-based generators, but often degrades perceptual quality in ways that are not captured by the reward proxy. We identify a simple structural signature of this drift: across three post-training methods (NFT, AWM, DPO), RL fine-tuning inflates the per-step velocity norm $\|v_\theta\|$ by $5\%$ to $15\%$ relative to the reference. A form of norm inflation has been studied in classifier-free guidance (CFG), where rescaling the velocity back to a reference norm at inference time can mitigate the resulting artifacts. However, this inference-time correction does not transfer cleanly to RL: rescaling $v_\theta$ to match $\|v_{\text{ref}}\|$ at inference time neither improves reward nor fixes the quality degradation, because the inflation is co-adapted into the model weights. Furthermore, an adjoint sensitivity analysis shows that velocity magnitude rescaling carries no coherent first-order reward signal at the batch level, indicating that suppressing norm inflation is unlikely to remove a consistently reward-carrying component. Since inference-time renormalization fails while norm suppression carries no reward cost, training-time intervention is the appropriate strategy. Together, these findings motivate NormGuard, a hinge penalty that activates only when $\|v_\theta\|$ exceeds $\|v_{\text{ref}}\|$ and composes additively with any velocity-local base loss. Across two base models, three post-training methods, and two reward proxies, NormGuard consistently improves MLLM-judged image quality and forensic realism while preserving reward, with gains that amplify under few-step inference and are not explained by early stopping.
Read the full abstract ▾
All authors Tianlin Pan, Lianyu Pang, Cheng Da, Huan Yang, Changqian Yu, Kun Gai, Wenhan Luo
Machine Learning
Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization
Clément Bonet, Pierre-Cyril Aubin-Frankowski and Youssef Mroueh
Optimizing functionals over the space of probability measures is now ubiquitous in machine learning. A widely used approach is to perform the optimization directly over the Wasserstein space, but many objective functionals of practical interest are non-convex along Wasserstein geodesics, making the analysis of standard first-order methods challenging. In this work, we study a class of objectives over the Wasserstein space that admit a difference-of-convex (DC) decomposition and we lift the classical convex-concave procedure (CCCP) to this setting. Under smoothness and strong convexity assumptions on the convex components of the decomposition, we prove almost stationarity along the iterates of the resulting
algorithm. Our main focus is on the Maximum Mean Discrepancy (MMD) and the Energy Distance (ED) functionals, for which we develop explicit Wasserstein DC decompositions, and establish local convergence of the scheme under mild assumptions. Empirically, we show that well-chosen DC decompositions yield faster and more stable convergence than Wasserstein gradient descent on these MMD objectives.
Read the full abstract ▾
Machine Learning
RS-Diffuser: Risk-Sensitive Diffusion Planning with Distributional Value Guidance
Shiqiang Gong
Offline reinforcement learning enables policy learning from fixed
datasets without additional environment interaction, making it appealing for safety-critical applications where online exploration is costly or unsafe.
Diffusion-based decision-making methods have recently achieved strong performance in offline RL by modeling rich,
multimodal trajectory distributions. However, existing
diffusion planners are typically risk-neutral and therefore may overlook rare but catastrophic outcomes that are crucial in real-world deployment. In this work, we propose RS-Diffuser, a risk-sensitive offline
diffusion planning framework that combines
diffusion-based trajectory generation with distributional value critics. RS-Diffuser learns a
diffusion planner over future state trajectories, a separate inverse dynamics model for action decoding, and a Monte Carlo distributional critic that estimates the full return distribution of candidate plans through
quantile regression. At sampling time, we incorporate a risk-sensitive guidance signal into the denoising process, using gradients computed from tail-aware objectives such as Conditional Value at Risk to steer generation toward desired risk profiles. As a result, a single trained model can flexibly produce risk-averse, risk-neutral, or risk-seeking behaviors by changing only the inference-time risk parameter. Extensive experiments on risk-sensitive D4RL and risky
robot navigation benchmarks demonstrate that RS-Diffuser achieves state-of-the-art performance, improving both overall return and worst-case robustness while reducing safety violations.
Read the full abstract ▾
Machine Learning
Layerwise Progressive Freezing: A Training Scaffold for Depth-Scalable Binary Networks
Evan Gibson Smith and Bashima Islam
Training binary neural networks (BNNs) from scratch is dominated by the straight-through
estimator (STE), whose forward/backward mismatch produces severe accuracy degradation as networks deepen. We study an orthogonal axis: when and where binarization is enforced during training. We introduce StoMPP (
Stochastic Masked Partial Progressive Binarization), which gradually replaces clipped weights and activations with their hard binary counterparts layer by layer from input to output, using
stochastic partial masks with soft refresh. StoMPP delivers two complementary benefits. As a standalone training rule, it provides a fully STE-free procedure that improves over vanilla STE with gains that grow with depth (ResNet-50 BNN: +18.0/+13.5/+3.8 on CIFAR-10/100/ImageNet), and the pattern holds across ResNet-18/34/50, MobileNetV2, and BERT fine-tuning. Composed with surrogate gradients by applying STE only to frozen entries, it reaches +27.1/+19.8/+17.7 over vanilla STE on the same setting. Underlying both regimes is a single mechanistic finding: progression order is decisive. Forward layerwise progression prevents depth collapse, reverse progression collapses to near-chance, and binary-weight networks (without binary activations) are insensitive to order. We trace this asymmetry to activation-induced gradient blockades: a committed binary activation severs gradient flow upstream, and ordering controls when these blockades form. To isolate the progression's contribution from any benefit conferred by STE, we conduct all ablations in the STE-free regime; the resulting characterization (schedule, refresh, ordering, dynamics) thus reflects the progression itself rather than its interaction with surrogate gradients.
Read the full abstract ▾
Machine Learning
PerturbCellRL: Aligning Distributions and Grounding Biology via Post-Training Perturbation Generators
Dongxia Wu, Mingyu Li, Yuhui Zhang et al. (+4 more)
Single-
cell perturbation models can reduce costly wet-lab screening by predicting how
cells respond transcriptionally to interventions. Recent advances in flow-matching have enabled population-level prediction of cellular responses. However, flow-matching training can fail to recover certain target distributions even within the model family, limiting its ability to capture cellular heterogeneity. We first prove that post-training can recover these distributions, then introduce PerturbCellRL, a reinforcement learning framework that post-trains single-
cell perturbation generators using per-
cell rewards. The central component is a gene-expression energy witness that translates population-level discrepancies into per-
cell feedback. We further prove that this reward's policy gradient points toward better distributional
alignment. Two complementary rewards, calibrated on real
cells, penalize atypical expression profiles and insufficient pathway-level responses to perturbations. Across genetic and chemical perturbation
benchmarks, PerturbCellRL substantially improves distributional
alignment and recovers pathway enrichment patterns more faithfully. These results establish reward-guided post-training as an effective strategy for improving both distributional accuracy and biological fidelity in perturbation prediction.
Read the full abstract ▾
All authors Dongxia Wu, Mingyu Li, Yuhui Zhang, Anurendra Kumar, Emma Lundberg, Serena Yeung-Levy, Emily B. Fox
Machine Learning
Flexformer: Flexible Linear Transformer with Learnable Attention Kernel
Haoran Zhang and Feng Zhou
Transformer models rely on
attention mechanism to capture long-range dependencies but suffer from quadratic complexity, limiting their scalability to long sequences. Kernel-based linear
attention reduces this complexity but typically relies on fixed or weakly learnable kernels, restricting expressiveness and performance. In this work, we propose Flexformer, a flexible linear
Transformer that learns
attention kernels in a fully data-driven manner. Flexformer builds on random
Fourier feature-based linear
attention and treats spectral frequencies as trainable parameters, enabling the model to learn a broad family of
attention kernels. We develop both stationary and nonstationary variants, with the latter offering strictly greater expressiveness. Extensive experiments on
language modeling and sequence classification demonstrate that Flexformer consistently outperforms baselines. Moreover, Flexformer can be effectively distilled from pretrained
Transformers to recover softmax
attention and exhibits strong kernel transferability across domains, achieving both high efficiency and competitive performance on long-sequence tasks.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.IR
End-to-End Dynamic Sparsity for Resource-Adaptive LLM Inference
Yuhang Chen, Jinhao Duan, Ruichen Zhang et al. (+11 more)
Large
Language Models (LLMs) inference is typically deployed under a static resource assumption, where models execute a fixed computational graph regardless of the runtime environment. However, real-world cloud infrastructure is inherently dynamic, characterized by fluctuating availability (e.g., spot instance preemption) and tiered Quality-of-Service requirements. In such volatile settings, static models are inflexible: they either crash under resource constraints or waste compute on redundant operations. To bridge this gap, we propose Learning to Allocate (L2A), an end-to-end framework for resource-adaptive inference. Unlike prior methods that condition only on input difficulty, we formulate inference as a constrained allocation problem conditioned on both the input and the runtime resource budget itself. We introduce lightweight, budget-conditioned and input-aware gating networks integrated into the LLM. These gates are trained via a unified objective that jointly optimizes task performance, logical consistency, and resource costs along three axes matching how real-world dynamics manifest: layer skipping for
memory and depth pressure, head
pruning for throughput contention, and
reasoning-
token reduction for latency tightening. This lets the model learn a budget-aware policy beyond input difficulty alone: it adaptively configures its computational footprint with respect to real-time resource dynamics, maximizing
reasoning depth when resources permit while enforcing strict frugality when budgets tighten. A single L2A model traces the entire compute-accuracy Pareto frontier on Llama-3-8B and Qwen-3-4B: at up to 34% realized layer sparsity, it stays within 0.6% of the dense baseline on GSM8K, with the same gap holding zero-shot on out-of-distribution tasks, while every static or heuristic baseline requires a separately tuned model and still drops by 5-10% at comparable inference time.
Read the full abstract ▾
All authors Yuhang Chen, Jinhao Duan, Ruichen Zhang, Mingfu Liang, Xiaohan Wei, Yunchen Pu, Fei Tian, Chonglin Sun, Parish Aggarwal, Frank Shyu, Luke Simon, Sandeep Pandey, Tianlong Chen, Xi Liu
Machine Learning
The Weakest Link Tells It All: Outcome-Supervised Process Reward Modeling via Learnable Credit Assignment
Tianyu Jia, Yue Fang, Hongxin Ding et al. (+6 more)
Process reward models (PRMs) enhance the
reasoning capabilities of large
language models (LLMs) by providing fine-grained feedback, yet training PRMs typically requires expensive stepwise annotations. Outcome-supervised PRMs offer a scalable alternative by learning from final-answer correctness alone, but this introduces a fundamental credit assignment challenge, i.e., attributing outcomes to responsible
reasoning steps. Existing approaches rely on either uniform or causal assignment, both of which fail to anchor credit in step correctness and thus hinder process error identification. In this work, we propose Outcome-Supervised Process Reward Modeling via Learnable Credit Assignment (LCA), an outcome-supervised PRM framework that jointly learns credit assignment and reward modeling under the principle of Weakest Link Assignment: a
reasoning chain is as strong as its weakest link. To address mutual dependence between credit assignment and reward modeling, we formalize outcome-supervised PRM as a Multiple Instance Learning (MIL) problem and introduce Softmax-Weighted-Sum (SWS) pooling, an MIL pooling technique tailored for strong dependence and redundancy among
reasoning states. We establish instance-level Bayes consistency of the resulting learning objective under mild assumptions. Extensive experiments demonstrate that LCA outperforms state-of-the-art outcome-supervised PRMs across multiple tasks and backbones. Code is available at https://github.com/JiaTianyu20031016/MIL.git.
Read the full abstract ▾
All authors Tianyu Jia, Yue Fang, Hongxin Ding, Rihong Qiu, Zhibang Yang, Zhijing Wu, Junfeng Zhao, Yasha Wang, Xu Chu