Machine Learning
A Tight Theory of Error Feedback Algorithms in Distributed Optimization
Daniel Berg Thomsen, Adrien Taylor and Aymeric Dieuleveut
Communication costs are a major bottleneck in distributed learning and first-order optimization. A common approach to alleviate this issue is to compress the gradient information exchanged between
agents. However, such compression typically degrades the convergence guarantees of gradient-based methods. Error feedback mechanisms provide a simple and computationally cheap remedy for this issue, but numerous variants have been proposed, and their relative performance remains poorly understood. This paper provides tight convergence analyses for two of the main error-feedback
algorithms from the literature, the classic Error Feedback method (EF) and Error Feedback 21 (EF21), by identifying optimal step-size choices and constructing optimal Lyapunov functions tailored to each method. The results hold independently of the number of
agents and recover the known best guarantees possible in the single-
agent regime.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
LongTraceRL: Learning Long-Context Reasoning from Search Agent Trajectories with Rubric Rewards
Nianyi Lin, Jiajie Zhang, Lei Hou et al. (+1 more)
Long-context
reasoning remains a central challenge for large
language models, which often fail to locate and integrate key information in extensive distracting content. Reinforcement learning with verifiable rewards (RLVR) has shown promise for this task, yet existing methods are limited by low-confusability distractors and sparse, outcome-only reward signals that cannot supervise intermediate
reasoning steps. To address these issues, we introduce LongTraceRL. For data construction, we generate multi-hop questions via knowledge graph random walks and leverage search
agent trajectories to build tiered distractors: documents the
agent read but did not cite (high confusability) and documents that appeared in search results but were never opened (low confusability), producing training contexts that are far more challenging than those built by random sampling or one-shot search. For reward design, we propose a rubric reward that uses the gold entities along each
reasoning chain as fine-grained, entity-level process supervision. This rubric reward is applied only to responses with correct final answers (positive-only strategy), distinguishing the
reasoning quality among correct responses and preventing reward hacking. Experiments on three
reasoning LLMs (4B–30B) across five long-context
benchmarks demonstrate that LongTraceRL consistently outperforms strong baselines and encourages comprehensive, evidence-grounded
reasoning. Codes,
datasets and models are available at \href{https://github.com/THU-KEG/LongTraceRL}{https://github.com/THU-KEG/LongTraceRL}.
Read the full abstract ▾
All authors Nianyi Lin, Jiajie Zhang, Lei Hou, Juanzi Li
Machine Learning
Giving Sensors a Voice: Multimodal JEPA for Semantic Time-Series Embeddings
Utsav Dutta, Gerardo Pastrana, Sina Khoshfetrat Pakazad et al. (+1 more)
Transformer-based architectures have advanced sequence modeling in
language and
vision, yet general-purpose representation learning for heterogeneous multivariate time series remains underexplored. We introduce CHARM (Channel-Aware Representation Model), which incorporates channel-level textual descriptions into a
Transformer encoder equivariant to channel order. CHARM is trained with a Joint
Embedding Predictive Architecture (JEPA) and a novel loss promoting informative, temporally stable
embeddings; latent-space prediction encourages robustness to sensor noise while description-aware gating provides interpretability through learned inter-channel relationships. Across anomaly detection, classification, and short- and long-term
forecasting, the learned
embeddings achieve strong performance using only a linear probe. Performance is driven primarily by the JEPA objective and conditioning architecture, with text descriptions serving as channel identifiers for cross-
dataset generalization.
Read the full abstract ▾
All authors Utsav Dutta, Gerardo Pastrana, Sina Khoshfetrat Pakazad, Henrik Ohlsson
Machine Learning
Effective Biological Representation Learning by Masking Gene Expression
Kian Kenyon-Dean, Alina Selega, Ihab Bendidi et al. (+5 more)
RNA sequencing produces rich and diverse
datasets of gene expression, offering compelling insights into cellular state and function that have many applications in drug discovery. Modeling such data is challenging due to inherent technical noise and experimental batch effects, as evidenced by many existing transcriptomic foundation models (FMs) underperforming relative to linear baselines. Such results raise the question of whether deep representation learning provides a distinct advantage over the direct use of raw transcript counts. Our work explores this by developing a new self-supervised model, TxFM, with a focus on inductive representation learning evaluations. TxFM employs a masked autoencoding approach tailored to diverse RNA-seq count data, and our ablation study empirically identifies crucial architecture configurations required for strong transfer performance. Additionally, we curate a public training corpus, DiverseRNA-1.4M, and find that TxFM trained on this curated
dataset yields high-fidelity gene representations that outperform FMs trained on atlas-scale corpora over 100x larger. Overall, our results indicate that inductive self-supervised learning is a viable modeling approach for transcriptomics representation, provided a careful synthesis of model architecture and training data curation.
Read the full abstract ▾
All authors Kian Kenyon-Dean, Alina Selega, Ihab Bendidi, Jordan M. Sorokin, Luca Bertinetto, David Errington, Hayley Donnella, Oren Kraus
Machine Learning
Functional Attention: From Pairwise Affinities to Functional Correspondences
Jiefang Xiao, Maolin Gao, Simon Weber et al. (+2 more)
Learning mappings between infinite-dimensional function spaces, or operator learning, is essential for many machine learning applications. Although
transformer-based operators are popular, they often rely on
token-wise
attention. These methods treat continuous fields as discrete
tokens and usually ignore the global functional structure. We introduce Functional
Attention, which reinterprets
attention as a functional correspondence between adaptive bases. Inspired by
geometric functional maps, our method replaces softmax affinities with structured linear operators. This yields a compact, generalizable, resolution-invariant representation that explicitly captures global dependencies. Experiments demonstrate that Functional
Attention can match state-of-the-art performance in many operator learning tasks, including solving PDEs, 3D
segmentation, and regression, while remaining robust to varying discretizations. Project page is available at https://github.com/xjffff/FUNCATTN.
Read the full abstract ▾
All authors Jiefang Xiao, Maolin Gao, Simon Weber, Guandao Yang, Daniel Cremers
Machine Learning
Positional versus Symbolic Attention Heads: Learning Dynamics, RoPE Geometry, and Length Generalization
Felipe Urrutia, Juan José Alegría, Cinthia Sanchez Macias et al. (+3 more)
Transformer-based
language models are widespread in today's society. As such, understanding the mechanisms by which they solve structured tasks and predicting how they may behave in novel scenarios is of great importance for safe deployment. We study the learning dynamics of
attention heads in a controlled setting by training a decoder-only
Transformer (GPT-J) on two structurally equivalent multi-hop
reasoning tasks: a number task requiring positional
reasoning and a letter task requiring symbolic
reasoning. Using a recently introduced metric that classifies
attention-head behavior as positional or symbolic for a given
prompt, we show that successful learning is associated with the emergence of pure heads, i.e., heads that express themselves as either positional or symbolic. Despite the tasks' structural equivalence, they impose different mechanistic demands: the number task requires both positional and symbolic heads, whereas the letter task requires only symbolic heads. We then identify the computational roles of these heads, characterize the basic functions they implement, and give theoretical constructions showing how single-layer RoPE-based
attention can realize these functions through geometrically interpretable query, key, and value operations. This analysis yields a quantitative separation between positional and symbolic mechanisms in their robustness to longer sequences, formalized through a novel notion of discrepancy. We empirically validate the resulting predictions in both controlled and real-world models, showing that symbolic mechanisms extrapolate more reliably to longer sequences while positional mechanisms face sharper limitations.
Read the full abstract ▾
All authors Felipe Urrutia, Juan José Alegría, Cinthia Sanchez Macias, Jorge Salas, Cristian B. Calderon, Cristobal Rojas
Machine Learning
The Dynamic-Probabilistic Consistency Gap in Chaotic Surrogate Modeling
Andre Herz, Matthijs Pals, Daniel Durstewitz et al. (+1 more)
Dynamical systems reconstruction (DSR) aims to learn surrogate models that capture the dynamics underlying time-series data. Reliably deploying these surrogates requires uncertainty estimates consistent with the learned dynamics. We expose a dynamic-probabilistic consistency (DPC) gap: the pursuit of finite-horizon probabilistic objectives can degrade dynamics or decouple predictive uncertainty from the local tangent dynamics it ought to reflect. We isolate three mechanisms behind this gap: core collapse, noise masking, and blind uncertainty. Specifically, we show that open-loop Gaussian rollout objectives can penalize Jacobian-generated covariance growth in chaotic systems, encouraging optimization shortcuts that weaken physical expansion or decouple uncertainty from it. To mitigate this gap, we propose KAFFEE (Kalman-Aware Framework For Ergodic Emulation), a differentiable extended Kalman filter-based training framework that evaluates likelihood on local predictive residuals (innovations) while transporting covariance through learned local Jacobians. On
stochastic hyperchaotic Lorenz-96, KAFFEE reduces the identified failure modes, improves reconstruction of dynamical invariants relative to open-loop objectives, and maintains competitive predictive scores. We further show that the DPC gap appears when probabilistically adapting a DSR foundation model across 13 chaotic systems, where KAFFEE enables in-context
Bayesian filtering while largely preserving zero-shot dynamics.
Read the full abstract ▾
All authors Andre Herz, Matthijs Pals, Daniel Durstewitz, Georgia Koppe
Machine Learning · Cross-listed from cs.CV
Automated Prediction of Postoperative Pancreatic Fistula Using Preoperative Computed Tomography
Ashok Choudhary, Chris Varghese, Leo Y. Li-Han et al. (+5 more)
Postoperative pancreatic fistula (POPF) is a serious complication after pancreatic resection, increasing morbidity, hospital stay, and healthcare costs. We present an automatic, end-to-end deep learning pipeline-from pancreatic
segmentation to classification-for preoperative POPF risk estimation and stratification using preoperative CT scans. A data set with auto-segmented pancreas volumes and
surgical outcomes was used to evaluate multiple architectures, including a custom lightweight 3D CNN baseline (CNN3D), R(2+1)D ResNet-18, and ResNet-MC3-18 models. Evaluation across multiple 3D architectures demonstrated promising predictive performance. This approach offers a clinically valuable tool and a methodological
benchmark for pancreas-specific CT classification, supporting improved preoperative decision-making in pancreatic surgery.
Read the full abstract ▾
All authors Ashok Choudhary, Chris Varghese, Leo Y. Li-Han, Frank G. Lee, Ellen L. Larson, Elizabeth B. Habermann, Cornelius A. Thiels, Hojjat Salehinejad
Machine Learning · Cross-listed from cs.CV
RayDer: Scalable Self-Supervised Novel View Synthesis from Real-World Video
Ulrich Prestel, Stefan Andreas Baumann, Nick Stracke et al. (+1 more)
Self-supervised novel view synthesis (NVS) remains challenging to scale, despite the abundance of
video data, largely due to the brittleness of training on realistic
videos and the hard-to-predict scaling behavior of multi-network system designs. We introduce RayDer, a unified, feed-forward
transformer that consolidates camera estimation, scene reconstruction, and rendering into a single backbone, turning self-supervised NVS into a well-posed single-model scaling problem. A minimal dynamic state, treated as a nuisance factor, absorbs time-varying content and enables stable training on unconstrained real-world
video. Importantly, RayDer keeps static-scene NVS as its target task: dynamic content is leveraged purely as scalable supervision, not reconstructed as in dynamic-scene (4D) NVS. Across multiple model sizes and orders of magnitude in data, RayDer exhibits clean power-law scaling with data and compute, and outperforms static-scene data mixtures. On a large number of
benchmarks, RayDer achieves strong zero-shot open-set performance competitive with state-of-the-art supervised approaches. Project Page: https://compvis.github.io/rayder
Read the full abstract ▾
All authors Ulrich Prestel, Stefan Andreas Baumann, Nick Stracke, Björn Ommer
Machine Learning · Cross-listed from cond-mat.soft
Discovering Thermodynamically Admissible Dissipation Potentials via Grammar-Based Symbolic Regression
Federico Califano and Jacopo Ciambella
Constitutive laws for inelastic materials must satisfy strict thermodynamic admissibility requirements, yet current data-driven approaches sacrifice interpretability, even when formal guarantees are provided by physics-encoded architectures. We propose a symbolic regression framework for the data-driven discovery of dissipation potentials governing the
evolution of internal variables within the Generalized Standard Materials (GSM) formalism. Starting from the Clausius–Duhem inequality, we enforce the thermodynamic requirements, convexity and non-negativity, that the dual dissipation potential must satisfy to guarantee non-negative mechanical dissipation. These requirements are formulated in the general subdifferential setting, encompassing rate-dependent (viscoelastic) and viscoplastic dissipative mechanisms, including potentials with genuine elastic domains, within a unified framework. Candidate potentials are generated by a composition-extended convexity-preserving grammar that guarantees thermodynamic admissibility by construction. The framework is validated on synthetic
datasets spanning Newtonian, power-law, and Bingham viscoplastic ground truths under process and measurement noise, and on experimental oscillatory shear measurements of a synthetic elastomer across multiple strain amplitudes and frequencies, where the discovered potentials reproduce the amplitude-dependent softening of the dynamic
moduli and outperform a calibrated linear Zener baseline.
Read the full abstract ▾
Machine Learning
Value Functions as Supermartingale Certificates
Alessandro Abate, Daniel Contro, Mirco Giacobbe et al. (+2 more)
Certification methods for
stochastic systems provide sufficient proof rules, based on real-valued supermartingale certificates, to determine the almost-sure satisfaction of $\omega$-regular properties (and therefore of linear temporal logic) over general state spaces, encompassing both countably infinite and continuous state spaces. Conversely, reinforcement learning (RL) methods for $\omega$-regular tasks have received considerable
attention, but they typically lack formal guarantees that the learned policy satisfies the specification, except possibly for finite state and action spaces. We bridge these two lines of research by establishing a novel theoretical connection: under an appropriate reward, the value function associated to a policy that almost surely satisfies an $\omega$-regular property encodes a Streett supermartingale certificate for that specification. Our results, validated experimentally on finite Markov decision processes, hold for finite, countably infinite, and continuous state spaces, suggesting a principled route to certificate synthesis via RL.
Read the full abstract ▾
All authors Alessandro Abate, Daniel Contro, Mirco Giacobbe, Agustín Martínez-Suñé, Diptarko Roy
Machine Learning
Chem-PerturBridge: a harmonized compendium of small molecule perturbation transcriptomic effects
Artur Szałata, Olga Novitskaia, Maiia Shulman et al. (+3 more)
Large perturbation models require training data encompassing chemical, cellular, and assay diversity. Current transcriptomic resources for small-molecule modeling, however, are fragmented across technologies, metadata conventions, controls, doses, and preprocessing pipelines. We introduce Chem-PerturBridge, a harmonized multi-
dataset resource comprising over 37k compounds, 136 cellular contexts, and 1.25M transcriptomic samples across eight assay types, with standardized identifiers, metadata, and replicate-aware condition-level effects. We use the resource to evaluate matched-condition agreement across
datasets and replicate agreement within
datasets. Matched same-compound conditions generally show weak agreement in fine-grained logFC rankings and magnitudes across most
dataset pairs, often falling below same-context different-compound baselines. In contrast, logFC direction agreement is substantially more stable and usually exceeds these baselines. We further evaluate Chem-PerturBridge as a pretraining resource for compound representation learning. Under a compound-held-out OP3 evaluation split,
embeddings pretrained on Chem-PerturBridge improve over L1000-only
embeddings, Morgan fingerprints, and the descriptor-free OP3 baseline across metrics. An extensive molecule-holdout evaluation across 11
datasets further shows that models trained on Chem-PerturBridge outperform or match those that are not. Chem-PerturBridge therefore supports both diagnostic evaluation of cross-
dataset signature agreement and model-oriented reuse of heterogeneous perturbation transcriptomic data.
Read the full abstract ▾
All authors Artur Szałata, Olga Novitskaia, Maiia Shulman, Matthew Mella, Altynbek Zhubanchaliyev, Fabian J. Theis
Machine Learning
On the Relationship Between Activation Outliers and Feature Death in Sparse Autoencoders
Elana Simon, Etowah Adams and James Zou
Sparse autoencoders (SAEs) decompose neural network activations into interpretable features, but many learned features never activate, a problem called feature death that wastes dictionary capacity and can reintroduce superposition. Death rates vary dramatically between models: near-zero on GPT-2, over 70% on AlphaFold3 with identical configurations. We find that dimension-level activation outliers (dimensions whose mean magnitude is large relative to per-
token variation) cause this by shifting pre-activations at initialization based on each feature's
alignment with the activation mean. Features anti-aligned with the mean receive permanently negative pre-activations and never fire. We formalize outlier severity as $\gamma = \|\mu\|/\|\sigma\|$; it predicts initial death rates (Spearman $\rho = 0.89$ for dead-by-TopK, $0.82$ for dead-by-ReLU) across 454 model-layer combinations spanning
language,
vision,
protein, and genomic models. Dead features can revive during training, but recovery requires the SAE bias to learn the activation mean, a process that is prohibitively slow at high $\gamma$. Mean-centering (subtracting the activation mean) sidesteps this and eliminates outlier-induced death across all tested models, confirming the mechanism and providing a principled basis for when and why this preprocessing step is necessary.
Read the full abstract ▾
Machine Learning
Skill Reuse as Compression in Agentic RL
Zhikun Xu, Yu Feng, Jacob Dineen et al. (+3 more)
Large
language model
agents trained with reinforcement learning (RL) often learn brittle, task-specific shortcuts. We hypothesize that
agents generalize better when their successful trajectories are structurally compressible, decomposed into a small set of reusable abstract patterns. To formalize this, we introduce ReuseRL, which grounds agentic RL in the Minimum Description Length (MDL) principle. ReuseRL extracts a shared skill dictionary from successful trajectories and augments the RL objective with a
segmentation cost, explicitly penalizing idiosyncratic behaviors that encode poorly. We prove a PAC-Bayes bound guaranteeing that a dictionary extracted from successful trajectories has bounded expected description length on future successful behavior. Across ALFWorld, TextWorld-Cooking, and Countdown-Stepwise, ReuseRL improves in- and out-of-distribution success over vanilla GRPO and strong round-length baselines.
Read the full abstract ▾
All authors Zhikun Xu, Yu Feng, Jacob Dineen, Taiwei Shi, Jieyu Zhao, Ben Zhou
Machine Learning
When Are Multimodal Predictions Biologically Supported? A Diagnostic Evaluation Framework
Dylan Steiner, Gustavo Arango-Argoty, Gerald Sun et al. (+1 more)
Multimodal models in oncology can produce accurate predictions, but accurate prediction does not reveal whether the model has learned biology that is shared across modalities, biology confined to one modality, or spurious correlations that reflect confounders rather than genuine biology. We introduce DECAT, a model-agnostic post-hoc evaluation framework that classifies
multimodal representations into four diagnostic scenarios for a given task and modality, using five null-referenced metrics and a rule-based decision procedure. The framework operates on learned representations, requires no knowledge of which specific confounder is present, and returns indeterminate when the evidence is insufficient. We validate DECAT on synthetic data across four
multimodal model classes (over 2,500 trained representations) and on real data from 8,979 TCGA patients, evaluating both
multimodal embeddings and five pretrained pathology foundation models. Entangled models (e.g., CLIP) achieve near-perfect shared biology detection but falsely claim shared biology in the majority of cases where it is absent on real foundation model
embeddings. This false claim rate increases with confound strength so that larger cohorts and stronger representations produce more confident but still incorrect diagnoses. Applied to both
multimodal TCGA
embeddings and five pathology foundation models without paired RNA, DECAT detects confounding invisible to AUROC without requiring the confounder labels, as confirmed by post-hoc stratification.
Read the full abstract ▾
All authors Dylan Steiner, Gustavo Arango-Argoty, Gerald Sun, Etai Jacob
Machine Learning · Cross-listed from cs.CV
How can embedding models bind concepts?
Arnas Uselis, Darina Koishigarina and Seong Joon Oh
Humans easily determine which color belongs to which shape in multi-object scenes, an ability known as concept binding.
Vision-
language embedding models such as CLIP struggle with binding: they recognize individual concepts but fail to represent which concepts form which objects. Although CLIP behaves like a bag-of-concepts model in cross-modal
retrieval, object information is recoverable from its image and text
embeddings separately. We study this tension through the binding function, which maps concepts to scene
embeddings. We find that scene
embeddings decompose additively into object representations, explaining why uni-modal probes can recover object information. However, CLIP's binding function is high-complexity, which likely prevents the image and text encoders from learning a shared binding mechanism that generalizes to unseen concept combinations. We then ask whether this limitation is fundamental. We show that it is not. In controlled
transformer models trained from scratch, binding generalization emerges with sufficient data coverage. These models learn low-complexity binding functions characterized by multiplicative interactions between concepts, enabling systematic generalization. Code is publicly available at https://github.com/oshapio/binding-concepts-complexity.
Read the full abstract ▾
Machine Learning
On Efficient Scaling of GNNs via IO-Aware Layers Implementations
Daria Fomina, Daniil Krasylnikov, Alexey Boykov et al. (+3 more)
Graph Neural Networks (GNNs) are bottlenecked by sparse, irregular
memory access. Popular frameworks such as DGL and PyTorch
Geometric support general message passing, but complex layers often materialize edge-wise intermediates, increasing
memory traffic and limiting scalability on large graphs. We take an I/O- and arithmetic-intensity–centric view and show that widely used layers fall into three kernel families: SpMM-based convolutions, reduction-based aggregations, and
attention-based layers (GATv2/Graph
Transformer). For each family, we develop GPU kernels that reduce data movement, improve locality, and remain robust across realistic graphs. We also study graph reordering and find that its impact depends on the kernel mapping: it benefits neighbor-parallel (gather-dominated) kernels more consistently than feature-parallel designs. Empirically, our fused
attention kernels reach up to $\textbf{3.9}\times$ speedup for Graph
Transformer (median $\textbf{1.6}\times$), with Tensor Core (block-sparse) variants up to $\textbf{7.3}\times$ on locally dense graphs; for GATv2 we reach up to $\textbf{8.5}\times$ speedup (median $\textbf{2.0}\times$) while reducing peak
memory by up to $\textbf{76}\times$ (median $\textbf{6}\times$). Our degree-aware reduction kernels achieve up to $\textbf{10}\times$ speedup (median $\textbf{2.6}\times$). For SpMM-based layers, properly cached cuSPARSE achieves up to $\textbf{8}\times$ speedup over DGL and outperforms evaluated custom baselines in the majority of evaluations. We release our implementations as drop-in replacements to support reproducible, hardware-aware GNN acceleration.
Read the full abstract ▾
All authors Daria Fomina, Daniil Krasylnikov, Alexey Boykov, Andrey Dolgovyazov, Vyacheslav Zhdanovskiy, Fedor Velikonivtsev
Machine Learning
Scalable Inference-Time Annealing with Surrogate Likelihood Estimators
Daniel Peñaherrera, Rishal Aggarwal and David Ryan Koes
A long standing challenge in computational chemistry and biophysics is efficiently sampling the Boltzmann distribution of molecules. Advances in generative modeling have been proposed to address the limitations of conventional sampling techniques by eliminating the computational cost of simulation. A promising direction is iteratively finetuning
diffusion models along a temperature ladder whereby training data is generated via importance sampling during inference-time annealing. Unfortunately, these methods require computing a divergence over the score field to estimate importance weights, rendering them intractable for larger systems. Here we present scalable inference-time annealing (SITA), which retrains flow-based models to generate samples at progressively lower temperatures using an energy-based model to facilitate fast surrogate likelihoods. We demonstrate state-of-the-art performance on both Alanine Dipeptide and Alanine Tripeptide while avoiding costly divergence terms. Our code is available at https://github.com/countrsignal/sita.git
Read the full abstract ▾
Machine Learning
Assign and Add: A Mechanistic Study of Compositional Arithmetic
Brady Exoo, Alberto Bietti and John Sous
Large
language models are able to compose skills in order to perform complex tasks, many of which might not have been seen during training. The details of how exactly this composition occurs remain elusive. In this paper, we study a mechanism for compositional generalization in
transformers by considering a simple controlled setting involving variable assignment and modular addition. By partitioning our training data into disjoint sets, we observe that small
transformers are able to generalize to previously unseen combinations of variables and numbers. Our mechanistic analysis shows that the same “modular addition” MLP module is used whether the inputs are given directly or indirectly through a separate variable assignment mechanism. We also analyze the training dynamics from an empirical lens, which reveals three phases of learning: first, modular addition is learned, then the structure required for variable assignment, and finally a refinement phase where the model generalizes to some hard sequences not seen in training. Finally, we provide a theoretical framework to explain how compositionality emerges from training dynamics. These results suggest that compositional generalization can be a natural consequence of the compositionality of internal mechanisms in~
transformers.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
Consolidating Rewarded Perturbations for LLM Post-Training
Zheyu Zhang, Shuo Yang and Gjergji Kasneci
Post-training of
language models is commonly framed as a sample-score-update loop implemented by gradient descent. A recent line of work, exemplified by RandOpt, relocates this loop to weight space, sampling Gaussian perturbations around a pretrained model and ensembling the top-K rewarded specialists at inference. While competitive with PPO and GRPO under matched training compute, this prediction-level ensemble incurs K forward passes per test example and does not extend cleanly to free-form generation. We ask whether the rewarded population can instead be folded into a single deployable model, replacing the inference-time ensemble with one consolidated update. A split-half analysis over 25 model-task pairs reveals reproducible low-rank structure in every case. We turn this geometry into CoRP (Consolidating Rewarded Perturbations), a gradient-free operator that combines reward-weighted aggregation, compatibility-aware reweighting, and a held-out validation gate, with no gradient flowing through the
language model. Across five
language models from 0.5B to 8B and five tasks covering math, code, and creative writing, CoRP improves the base model by 8.1 points on average. Using one tenth of RandOpt's perturbation budget, CoRP exceeds single-inference RandOpt by 6.5 points and recovers more than half of the gain of the 50-pass majority-vote ensemble, at one forward pass per test example.
Read the full abstract ▾
Machine Learning
Graphical einops: bridging tensor networks and computation graphs
Vincent Wang-Maścianica and Nikhil Khatri
Architecture diagrams are ubiquitous in deep learning, but they are usually only representational: the tensor-program identities they suggest are still proved by prose and tensor-axis manipulation. We introduce a formal graphical calculus for the structural fragment of tensor programming underlying einops, making such diagrams proof-enabling. Our calculus represents tensor axes as nested graded tubes around a base type. The tube boundary recovers the undirected tensor-network view of axes, while the directed interior retains the operational reading of computation graphs. The key rewrite is grade-naturality: sliding spectacles over tubes. Standard equivariance proofs become short diagrammatic derivations. We additionally demonstrate how our rewrite system may be applied to convert
attention masks into pre-processing operations, recovering efficient implementations of sparse
attention blocks.
Read the full abstract ▾
Machine Learning
Balanced LoRA: Removing Parameter Invariance to Accelerate Convergence
Valérie Castin, Kimia Nadjahi, Pierre Ablin et al. (+1 more)
Low-Rank Adaptation (LoRA) is the most widely adopted method for fine-tuning large
language models. Notably, LoRA is inherently overparameterized: multiple pairs of low-rank factors can yield the same adapted weight
matrix. We show–both theoretically and empirically–that these pairs exhibit significantly different condition numbers. As a result, converging to different loss minimizers directly impacts the convergence rate of LoRA. Building on this observation, we introduce Balanced Low-Rank Adaptation (BaLoRA), a variant of LoRA that projects iterates onto a balanced
manifold. This
manifold improves the conditioning of the loss landscape while preserving the adapted
matrix. The projection step is computationally lightweight and integrates seamlessly into existing fine-tuning pipelines. Empirically, BaLoRA converges faster than standard LoRA and achieves superior performance across a range of fine-tuning tasks.
Read the full abstract ▾
All authors Valérie Castin, Kimia Nadjahi, Pierre Ablin, Gabriel Peyré
Machine Learning
GPU Forecasters: Language Models as Selective Surrogates for Kernel Runtime Optimization
Zaid Khan, Justin Chih-Yao Chen, Jaemin Cho et al. (+2 more)
GPU kernels are the workhorse of modern deep learning, and optimizing them (via evolutionary search or coding
agents) usually requires repeated measurement on target hardware. While these measurements provide the ground-truth signal necessary for kernel search, they are costly, because each evaluation of a kernel requires compilation and repeated execution on a GPU. As improvements in LLM inference reduce the cost of writing novel kernels and LLM-driven searches scale to large search budgets, on-device evaluation becomes a bottleneck. To address this, we study how LLMs can serve as selective GPU surrogates for kernel evaluation, by
forecasting the performance of proposed kernels. A useful surrogate should be accurate, and it should be selective, by knowing when it could be wrong, and deferring to the GPU. To evaluate surrogates, we measure whether their forecasts are accurate, calibrated, and practically useful for recovering fast kernels under limited GPU-measurement budgets. Next, we study whether reinforcement learning can improve forecast accuracy and confidence calibration. Our experiments demonstrate that LLMs can accurately forecast relative kernel performance, that their utility can be improved through reinforcement learning. Used inside a kernel search, the surrogate lets the search consider several times as many candidates under the same GPU evaluation budget, and that leads to finding faster kernels than an equal-budget baseline. These results suggest that LLMs can play a broader role in kernel optimization, by acting as virtual models of a GPU rather than solely as kernel generators for search.
Read the full abstract ▾
All authors Zaid Khan, Justin Chih-Yao Chen, Jaemin Cho, Elias Stengel-Eskin, Mohit Bansal
Machine Learning
PithTrain: A Compact and Agent-Native MoE Training System
Ruihang Lai, Hao Kang, Haozhan Tang et al. (+6 more)
Mixture-of-Experts (MoE) has become the dominant architecture for frontier
language models. To meet this demand, production frameworks have built optimized MoE training stacks over years of engineering effort. Yet evolving these stacks for new architectures and system optimizations remains expensive. With the rise of AI coding
agents, they could automate parts of training-framework development and accelerate this
evolution. But applying them to these existing frameworks carries hidden costs, invisible to today's throughput-only evaluations. We name this missing dimension
agent-task efficiency (ATE): the cost of using coding
agents to understand, operate, and extend a framework. Grounded in four
agent-native design principles, we build PithTrain, a compact,
agent-native MoE training framework. We further introduce ATE-Bench, covering real-world training-framework tasks. Our evaluation shows PithTrain matches the throughput of production frameworks, and on ATE-Bench, PithTrain enables higher
agent-task efficiency, with up to 62% fewer
Agent Turns and 64% less Active GPU Time.
Read the full abstract ▾
All authors Ruihang Lai, Hao Kang, Haozhan Tang, Akaash R. Parthasarathy, Zichun Yu, Junru Shao, Todd C. Mowry, Chenyan Xiong, Tianqi Chen
Machine Learning
DRIFT: Decoupled Rollouts and Importance-Weighted Fine-Tuning for Efficient Multi-Turn Optimization
Jian Mu, Tianyi Lin, Chengwei Qin et al. (+2 more)
Large
language models are increasingly deployed in multi-turn interactive settings where users or environments can iteratively provide lightweight feedback. Unfortunately, optimizing such behavior presents a sharp dilemma in practice: online reinforcement learning is able to effectively address multi-turn dynamics but is prohibitively expensive due to the cost of generating full correction trajectories at every update, whereas offline supervised fine-tuning (SFT) is efficient but suffers from distribution shift and behavioral collapse. To this end, we novelly propose DRIFT (Decoupled Rollouts and Importance-Weighted Fine-Tuning), a framework that operationalizes the theoretical insight that the KL-regularized RL objective is equivalent to importance-weighted supervised learning. DRIFT decouples rollout from optimization by sampling offline interaction trajectories from a fixed reference policy, deriving return-based importance weights, and optimizing the policy via weighted SFT on the resulting
dataset. Empirically, we demonstrate that DRIFT matches or exceeds the performance of multi-turn reinforcement learning baselines while maintaining the training efficiency and simplicity of standard supervised fine-tuning. Code is available at https://github.com/2020-qqtcg/DRIFT.
Read the full abstract ▾
All authors Jian Mu, Tianyi Lin, Chengwei Qin, Zhongxiang Dai, Yao Shu
Machine Learning · Cross-listed from cs.GT
Used Car Salesbots? Honesty and Credulity of LLMs as Bargaining Agents under Partial Information
Antonio Valerio Miceli-Barone, Vaishak Belle and Shay B. Cohen
In this work we study
agents in simulated bargaining scenarios, where a buyer and a seller communicate through a text channel and attempt to negotiate mutually beneficial trades, under different information regimes (complete information, information asymmetry or mutual uncertainty). We evaluate their performance w.r.t. game-theoretical solutions and further investigate their honesty (their tendency to disclose or withhold information or to mislead and deceive) as well as their credulity (their tendency to trust or distrust information provided by the other
agent). We study zero-shot LLM
agents with simple prompting scaffolding as well as fine-tuned
agents, in order to investigate whether optimising the
agents to maximise financial profits makes them stronger negotiators but also more dishonest and less trusting. We find that off-the-shelf LLMs all substantially deviate from game-theoretical
equilibria, they attempt to lie about their private information but cannot efficiently exploit information asymmetries. Fine-tuning on financial utility makes the
agents stronger at achieving better deals but also more dishonest, highlighting the risks that optimising
agents for a task can have on their safety. We release our code and a
dataset of bargaining scenarios.
Read the full abstract ▾
Machine Learning · Cross-listed from stat.ME
Modeling Covariate Transition for Efficient Estimation of Longitudinal Treatment Effects in Randomized Experiments
Naoki Chihara, Tatsushi Oka, Yasuko Matsubara et al. (+2 more)
We present a regression-adjustment framework designed for the estimation of longitudinal treatment effects in randomized experiments under static regimes. While regression-adjustment methods are useful for variance reduction in randomized experiments by using pre-treatment
covariates, they usually focus only on average effects, from which we cannot obtain valuable insights into when the effects appear and how long they continue. To address this issue, we consider intermediate outcomes and evolving post-treatment
covariates over time, and we represent such dynamic trajectories using transition kernels. Furthermore, we establish the asymptotic normality and the semiparametric efficiency bound for our
estimator, enabling more powerful statistical inference. Simulation studies and empirical analysis using A/B test data from a streaming platform in Japan show the practical advantages of our method.
Read the full abstract ▾
All authors Naoki Chihara, Tatsushi Oka, Yasuko Matsubara, Yasushi Sakurai, Shota Yasui
Machine Learning
Flow map learning in nonlinear vector autoregressive models: influence of the feature-library structure on the training error
Markus Gross
Time series
forecasting often requires learning nonlinear and time-delayed dependencies. A paradigmatic class of
forecasting models are nonlinear vector autoregressive processes (NVAR), also known as next-generation reservoir computers (NG-RCs). These models approximate the Koopman operator on the space spanned by their explicit feature library. We consider the identifiability problem for learning Markovian nonlinear dynamical systems and show that the training error as a function of time resolution follows characteristic (pre-)asymptotic scaling laws. These laws depend on whether the feature library can represent the early Lie-series coefficients of the flow map (propagator) exactly or merely approximately. For dynamical systems governed by
polynomial vector fields, we demonstrate the mechanism for NVAR/NG-RC models with monomial and
Fourier feature libraries. We determine the dependence of the training error on the temporal resolution, the involved nonlinear degree, and the number of delay terms. While delay terms reduce the optimal one-step training error, they improve long-horizon forecasts only when the library provides sufficient nonlinearity. Thus, small training error coexists with weak generalization as the model class is mismatched to the true data-generating process. Numerical experiments on various chaotic dynamical systems confirm the theoretical predictions.
Read the full abstract ▾
Machine Learning
DG-CoLearn: An Efficient Collaborative Learning Framework for Dynamic Graphs
Ashley Hoi-Ting Au, Zikun Zhang, Ligang He et al. (+1 more)
Dynamic graph learning (DGL) is essential for modelling evolving graph data, but existing methods suffer from significant computational overhead due to repeated full-snapshot retraining and are not well-suited for collaborative settings with partitioned data. In realistic graph systems, cross-partition edges are unavoidable, but direct sharing of graph structure between clients may violate
privacy constraints. We propose DG-CoLearn, a client-oblivious collaborative dynamic graph learning framework built on incremental graph snapshot processing, which focuses computation on graph regions affected by temporal updates while preserving historical information through temporal modelling. This incremental design is consistently applied across the entire graph processing pipeline, including a server-mediated
embedding exchange mechanism to enable accurate multi-hop message passing without exposing raw cross-client structural information. Extensive experiments demonstrate that DG-CoLearn achieves up to 33.8$\times$ speedup in training time and 27.4$\times$ reduction in communication overhead, while consistently improving predictive performance on both node classification (up to 13.36% F1 improvement) and link prediction (up to 8.27% MAP improvement) tasks. These results highlight the effectiveness of DG-CoLearn in bridging efficiency, scalability, and client-to-client structural
privacy in collaborative dynamic graph learning.
Read the full abstract ▾
All authors Ashley Hoi-Ting Au, Zikun Zhang, Ligang He, Qiang Ni
Machine Learning
Jingwen Liu, Alexandr Andoni and Daniel Hsu
We introduce universal
transformers: fixed
transformers that can simulate any
transformer in a given class via a suitable input
embedding. Analogous to a universal Turing machine, the input
embedding encodes a description of the target model while all internal parameters remain fixed. We provide explicit sparse constructions achieving universality when the
embedding dimension is sufficiently large, and further show that universality is generic: randomly initialized
transformers are universal almost surely, which aligns with recent empirical results of Zhong and Andreas (2024). We empirically validate our theory on the algorithmic tasks of parenthesis balancing and multi-hop
reasoning. Our results suggest that much of a
transformer's expressive power may reside in its input representation rather than its learned weights.
Read the full abstract ▾