Machine Learning
Beyond GRPO and On-Policy Distillation: An Empirical Sparse-to-Dense Reward Principle for Language-Model Post-Training
Yuanda Xu, Hejian Sang, Zhengze Zhou et al. (+3 more)
In settings where labeled verifiable training data is the binding constraint, each checked example should be allocated to the model and reward density where it is most informative. We identify a reward-density principle: sparse sequence-level reward is most useful on models that can explore and discover better behavior, while dense
token-level teacher supervision is better suited for compressing that behavior into a smaller deployment model; accordingly, scarce labeled data should be used upstream on the strongest available teacher, then transferred downstream as dense supervision. We evaluate this rule through a four-stage workflow:teacher RL, forward-KL warmup, on-policy
distillation, optional post-bridge student RL on verifiable math and code generation with Qwen3, Llama, and Qwen2.5 models. At fixed Qwen3-1.7B deployment-student size, an RL-improved 8B teacher distilled through the dense bridge outperforms direct GRPO on the same student ($79.3\%$ vs.\ $75.9\%$ on MATH; $25.2\%$ vs.\ $19.8\%$ on AIME~2024), while transfer from the same teacher before RL underperforms. A component ablation confirms that each stage is load-bearing: the RL-improved teacher adds $7.8$ absolute points on MATH over a raw teacher of the same size, the forward-KL warmup adds $1.7$, and on-policy
distillation adds $3.3$. The teacher-quality ordering: raw-teacher transfer $<$ direct GRPO $<$ RL-teacher transfer, replicates on Llama-3.1-8B-Instruct with a Llama-3.3-70B-Instruct teacher and on Qwen2.5-7B-Instruct for code generation with a Qwen2.5-32B-Instruct teacher. The operational lesson is to avoid spending scarce labeled data on the least prepared policy: use sparse reward for teacher-side discovery, dense transfer for student compression, and student-side sparse reward only after the bridge. The OPD implementation used in our experiments is available at https://github.com/HJSang/OPSD_OnPolicyDistillation.
Read the full abstract ▾
All authors Yuanda Xu, Hejian Sang, Zhengze Zhou, Ran He, Zhipeng Wang, Alborz Geramifard
Machine Learning
MEME: Multi-entity & Evolving Memory Evaluation
Seokwon Jung, Alexander Rubinstein, Arnas Uselis et al. (+2 more)
LLM-based
agents increasingly operate in persistent environments where they must store, update, and reason over information across many sessions. While prior
benchmarks evaluate only single-entity updates, MEME defines six tasks spanning the full space defined by the multi-entity and evolving axes, including three not scored by prior work: Cascade and Absence (dependency
reasoning) and Deletion (post-removal state). Evaluating six
memory systems spanning three
memory paradigms on 100 controlled episodes, we find that all systems collapse on dependency
reasoning under the default configuration (Cascade: 3%, Absence: 1% in average accuracy) despite adequate static
retrieval performance.
Prompt optimization, deeper
retrieval, reduced filler noise, and most stronger LLMs fail to close this gap. Only a file-based
agent paired with Claude Opus 4.7 as its internal LLM partially closes the gap, but at ~70x the baseline cost, indicating closure currently depends on configurations that are not practical at scale. Code and data are available on the project page: https://seokwonjung-jay.github.io/meme-eval/.
Read the full abstract ▾
All authors Seokwon Jung, Alexander Rubinstein, Arnas Uselis, Sangdoo Yun, Seong Joon Oh
Machine Learning
Routers Learn the Geometry of Their Experts: Geometric Coupling in Sparse Mixture-of-Experts
Sagi Ahrac, Noya Hochwald and Mor Geva
Sparse Mixture-of-Experts (SMoE) models enable scaling
language models efficiently, but training them remains challenging, as routing can collapse onto few experts and auxiliary load-balancing losses can reduce specialization. Motivated by these hurdles, we study how routing decisions in SMoEs are formed mechanistically. First, we reveal a
geometric coupling between routers and their corresponding experts. For a given
token, the router weights for the selected expert and the expert weights processing it receive gradients along the same input direction, differing only in scalar coefficients. Thus, matched router–expert directions accumulate the same routed
token history. This theoretical coupling also appears empirically in routing dynamics. In a $1$B SMoE trained from scratch, higher router scores predict stronger expert
neuron activations, showing that routing decisions are mirrored inside the selected expert. Next, we analyze the effects of auxiliary load balancing on the router–expert
geometric coupling, showing that such losses break this structure by spreading input-directed gradients across router weights, making distinct router directions nearly three times more similar to each other. Last, we demonstrate the centrality of
geometric coupling for effective routing with a parameter-free online K-Means router, in which each expert maintains a running average of the hidden states routed to it and
tokens are assigned based on cosine similarity. Compared with auxiliary-loss and loss-free balancing, this router achieves the lowest load imbalance with only a modest perplexity increase, indicating that
geometric coupling captures a substantial part of what the router learns. Overall, our results explain how routers form assignment geometry that supports an effective division of labor.
Read the full abstract ▾
Machine Learning
KV-Fold: One-Step KV-Cache Recurrence for Long-Context Inference
Alireza Nadali, Patrick Cooper, Ashutosh Trivedi et al. (+1 more)
We introduce KV-Fold, a simple, training-free long-context inference protocol that treats the key-value (KV) cache as the accumulator in a left fold over sequence chunks. At each step, the model processes the next chunk conditioned on the accumulated cache, appends the newly produced keys and values, and passes the enlarged cache forward; the same one-step update is applied repeatedly, analogous to foldl in functional programming. Building on the KV cache concatenation primitive introduced for latent multi-
agent communication, we repurpose it as a chunk-to-chunk recurrence for long-context inference. When processing chunk t, the model attends to the KV cache carried from earlier chunks as a prefix, reusing its internal state across segments without modifying or retraining the model. Despite its simplicity, the induced recurrence is stable: per-step drift rises briefly and then saturates into a flat plateau that persists across deep chains. This plateau is insensitive to a 10,000x change in numerical precision, robust across chunk sizes, and consistent across model families. At the task level, KV-Fold preserves exact information over long distances. On a needle-in-a-haystack
benchmark, it achieves 100% exact-match
retrieval across 152 trials spanning contexts from 16K to 128K
tokens and chain depths up to 511 on Llama-3.1-8B, while remaining within the
memory limits of a single 40GB GPU. Compared to streaming methods, which trade fidelity for bounded
memory, KV-Fold maintains long-range
retrieval while operating as a sequence of tractable forward passes. Overall, our results show that frozen pretrained
transformers already support a stable form of KV-cache recurrence, providing a practical route to long-context inference without architectural changes or training.
Read the full abstract ▾
All authors Alireza Nadali, Patrick Cooper, Ashutosh Trivedi, Alvaro Velasquez
Machine Learning
Solve the Loop: Attractor Models for Language and Reasoning
Jacob Fein-Ashley and Paria Rashidinejad
Looped
Transformers offer a promising alternative to purely feed-forward computation by iteratively refining latent representations, improving
language modeling and
reasoning. Yet recurrent architectures remain unstable to train, costly to optimize and deploy, and constrained to small, fixed recurrence depths. We introduce Attractor Models, in which a backbone module first proposes output
embeddings, then an attractor module refines them by solving for the fixed point, with gradients obtained through implicit differentiation. Thus, training
memory remains constant in effective depth, and iterations are chosen adaptively by convergence. Empirically, Attractor Models outperform existing models across two regimes, large-scale
language-model pretraining and
reasoning with tiny models. In
language modeling, Attractor Models deliver a Pareto improvement over standard
Transformers and stable looped models across sizes, improving perplexity by up to 46.6% and downstream accuracy by up to 19.7% while reducing training cost. Notably, a 770M Attractor Model outperforms a 1.3B
Transformer trained on twice as many
tokens. On challenging
reasoning tasks, we show that our model with only 27M parameters and approximately 1000 examples achieves 91.4% accuracy on Sudoku-Extreme and 93.1% on Maze-Hard, scaling favorably where frontier models like Claude and GPT o3, fail completely, and specialized recursive reasoners collapse at larger sizes. Lastly, we show that Attractor Models exhibit a novel phenomenon, which we call
equilibrium internalization: fixed-point training places the model's initial output
embedding near
equilibrium, allowing the solver to be removed at inference time with little degradation. Together, these results suggest that Attractor Models make iterative refinement scalable by turning recurrence into a computation the model can learn to internalize.
Read the full abstract ▾
Machine Learning
High-arity Sample Compression
Leonardo N. Coregliano and William Opich
Recently, a series of works have started studying variations of concepts from learning theory for product spaces, which can be collected under the name high-arity learning theory. In this work, we consider a high-arity variant of sample compression schemes and we prove that the existence of a high-arity sample compression scheme of non-trivial quality implies high-arity PAC learnability.
Read the full abstract ▾
Machine Learning
Search Your Block Floating Point Scales!
Tanmaey Gupta, Hayden Prairie, Xiaoxia Wu et al. (+10 more)
Quantization has emerged as a standard technique for accelerating inference for generative models by enabling faster low-precision computations and reduced
memory transfers. Recently, GPU accelerators have added first-class support for microscaling Block Floating Point (BFP) formats. Standard BFP
algorithms use a fixed scale based on the maximum magnitude of the block. We observe that this scale choice can be suboptimal with respect to quantization errors. In this work, we propose ScaleSearch, an alternative strategy for selecting these scale factors: using a fine-grained search leveraging the mantissa bits in microscaling formats to minimize the quantization error for the given distribution. ScaleSearch can be integrated with existing quantization methods such as Post Training Quantization and low-precision
attention, and is shown to improve their performance. Additionally, we introduce ScaleSearchAttention, an accelerated NVFP4-based
attention algorithm, which uses ScaleSearch and adapted prior techniques to ensure near-0 performance loss for causal
language modeling. Experiments show that ScaleSearch reduces quantization error by 27% for NVFP4 and improves
language model PTQ by up to 15 points for MATH500 (Qwen3-8B), while ScaleSearchAttention improves Wikitext-2 PPL by upto 0.77 points for Llama 3.1 70B. The proposed methods closely match baseline performance while providing quantization accuracy improvements.
Read the full abstract ▾
All authors Tanmaey Gupta, Hayden Prairie, Xiaoxia Wu, Reyna Abhyankar, Qingyang Wu, Austin Silveria, Pragaash Ponnusamy, Jue Wang, Ben Athiwaratkun, Leon Song, Tri Dao, Daniel Y. Fu, Chris De Sa
Machine Learning · Cross-listed from cs.AI
Towards Affordable Energy: A Gymnasium Environment for Electric Utility Demand-Response Programs
Jose E. Aguilar Escamilla, Lingdong Zhou, Xiangqi Zhu et al. (+1 more)
Extreme weather and volatile wholesale electricity
markets expose residential consumers to catastrophic financial risks, yet demand response at the distribution level remains an underutilized tool for grid flexibility and energy affordability. While a demand-response program can shield consumers by issuing financial credits during high-price periods, optimizing this sequential decision-making process presents a unique challenge for reinforcement learning despite the plentiful offline historical smart meter and wholesale pricing data available publicly. Offline historical data fails to capture the dynamic, interactive feedback loop between an electric utility's pricing signals and customer acceptance and adaptation to a demand-response program. To address this, we introduce DR-Gym, an open-source, online Gymnasium-compatible environment designed to train and evaluate demand-response from the electric utility's perspective. Unlike existing device-level energy simulators, our environment focuses on the
market-level electric utility setting and provides a rich observational space relevant to the electric utility. The simulator additionally features a regime-switching wholesale price model calibrated to real-world extreme events, alongside physics-based building demand profiles. For our learning signal, we use a configurable, multi-objective reward function for specifying diverse learning objectives. We demonstrate through baseline strategies and data snapshots the capability of our simulator to create realistic and learnable environments.
Read the full abstract ▾
All authors Jose E. Aguilar Escamilla, Lingdong Zhou, Xiangqi Zhu, Huazheng Wang
Machine Learning · Cross-listed from math.ST
A proximal gradient algorithm for composite log-concave sampling
Linghai Liu and Sinho Chewi
We propose an
algorithm to sample from composite log-concave distributions over $\mathbb{R}^d$, i.e., densities of the form $\pi\propto e^{-f-g}$, assuming access to gradient evaluations of $\alpha_f$-convex $f$ and a restricted Gaussian oracle (RGO) for $\alpha_g$-convex $g$, where $\alpha_g \ge 0$. The latter requirement means that we can easily sample from the density $\text{RGO}_{g,h,y}(x) \propto \exp(-g(x) -\frac{1}{2h}||y-x||^2)$, which is the sampling analogue of the proximal operator for $g$. If $f + g$ is $\alpha$-convex and $f$ is $\beta$-smooth, our sampler achieves $\varepsilon$ error in total variation distance using $\widetilde{\mathcal O}(\kappa \sqrt d \log^4(1/\varepsilon))$ oracle calls where $\kappa := 1 \vee \beta/\alpha$, with $0\le \alpha_g \le \beta$ and $\alpha := \alpha_f + \alpha_g > 0$, which matches prior state-of-the-art results for the case $g=0$. We further extend our results to cases where (1) $\pi$ is non-log-concave but satisfies a Poincaré or log-Sobolev inequality, and (2) $f$ is non-smooth but Lipschitz.
Read the full abstract ▾
Machine Learning
Multi-Stream LLMs: Unblocking Language Models with Parallel Streams of Thoughts, Inputs and Outputs
Guinan Su, Yanwu Yang, Xueyan Li et al. (+1 more)
The continued improvements in
language model capability have unlocked their widespread use as drivers of autonomous
agents, for example in coding or computer use applications. However, the core of these systems has not changed much since early instruction-tuned models like ChatGPT. Even advanced AI
agents function on message exchange formats, successively exchanging messages with users, systems, with itself (i.e. chain-of-thought) and tools in a single stream of computation. This bottleneck to a single stream in chat models leads to a number of limitations: the
agent cannot act (generate output) while reading, and in reverse, cannot react to new information while writing. Similarly, the
agent cannot act while thinking and cannot think while reading or acting on information. In this work, we show that models can be unblocked by switching from instruction-tuning for sequential message formats to instruction-tuning for multiple, parallel streams of computation, splitting each role into a separate stream. Every forward pass of the
language model then simultaneously reads from multiple input streams and generates
tokens in multiple output streams, all of which causally depend on earlier timesteps. We argue that this data-driven change remedies a number of usability limitations as outlined above, improves model efficiency through parallelization, improves model security through better separation of concerns and can further improve model monitorability.
Read the full abstract ▾
All authors Guinan Su, Yanwu Yang, Xueyan Li, Jonas Geiping
Machine Learning · Cross-listed from cs.CR
TextSeal: A Localized LLM Watermark for Provenance & Distillation Protection
Tom Sander, Hongyan Chang, Tomáš Souček et al. (+10 more)
We introduce TextSeal, a state-of-the-art watermark for large
language models. Building on Gumbel-max sampling, TextSeal introduces dual-key generation to restore output diversity, along with
entropy-weighted scoring and multi-region localization for improved detection. It supports serving optimizations such as speculative decoding and multi-
token prediction, and does not add any inference overhead. TextSeal strictly dominates baselines like SynthID-text in detection strength and is robust to dilution, maintaining confident localized detection even in heavily mixed human/AI documents. The scheme is theoretically distortion-free, and evaluation across
reasoning benchmarks confirms that it preserves downstream performance; while a
multilingual human evaluation (6000 A/B comparisons, 5
languages) shows no perceptible quality difference. Beyond its use for provenance detection, TextSeal is also “radioactive”: its watermark signal transfers through model
distillation, enabling detection of unauthorized use.
Read the full abstract ▾
All authors Tom Sander, Hongyan Chang, Tomáš Souček, Tuan Tran, Valeriu Lacatusu, Sylvestre-Alvise Rebuffi, Alexandre Mourachko, Surya Parimi, Christophe Ropers, Rashel Moritz, Vanessa Stark, Hady Elsahar, Pierre Fernandez
Machine Learning · Cross-listed from eess.SP
Enabling AI-Native Mobility in 6G: A Real-World Dataset for Handover, Beam Management, and Timing Advance
Mannam Veera Narayana, Rohit Singh, Deepa M.R et al. (+1 more)
To address the issues of high interruption time and measurement report overhead under user equipment (UE) mobility especially in high speed 5G use cases the use of AI/ML techniques (AI/ML beam management and mobility procedures) have been proposed. These techniques rely heavily on data that are most often simulated for various scenarios and do not accurately reflect real deployment behavior or user traffic patterns. Therefore, there is an utmost need for realistic
datasets under various conditions. This work presents a
dataset collected from a commercially deployed network across various modes of mobility (pedestrian, bike, car, bus, and train) and at multiple speeds to depict real time UE mobility. When collecting the
dataset, we focused primarily on handover (HO) scenarios, with the aim of reducing the HO interruption time and maintaining continuous throughput during and immediately after HO execution. To support this research, the
dataset includes timing advance (TA) measurements at various signaling events such as RACH trigger, MAC CE, and PDCCH grant which are typically missing in existing works. We cover a detailed description of the creation of the
dataset; experimental setup, data acquisition, and extraction. We also cover an exploratory analysis of the data, with a primary focus on mobility, beam management, and TA. We discuss multiple use cases in which the proposed
dataset can facilitate understanding of the inference of the AI/ML model. One such use case is to train and evaluate various AI/ML models for TA prediction.
Read the full abstract ▾
All authors Mannam Veera Narayana, Rohit Singh, Deepa M.R, Radha Krishna Ganti
Machine Learning
ORCE: Order-Aware Alignment of Verbalized Confidence in Large Language Models
Chen Li, Xiaoling Hu, Songzhu Zheng et al. (+2 more)
Large
language models (LLMs) often produce answers with high certainty even when they are incorrect, making reliable confidence estimation essential for deployment in real-world scenarios. Verbalized confidence, where models explicitly state their confidence in natural
language, provides a flexible and user-facing uncertainty signal that can be applied even when
token logits are unavailable. However, existing verbalized-confidence methods often optimize answer generation and confidence generation jointly, which can cause confidence-
alignment objectives to interfere with answer accuracy. In this work, we propose a decoupled and order-aware framework for verbalized confidence calibration. Our method first generates an answer and then estimates confidence conditioned on the fixed question–answer pair, allowing confidence optimization without directly perturbing the answer-generation process. To align confidence with correctness likelihood, we construct a sampling-based surrogate from multiple model completions and optimize rank-based reinforcement learning objectives that encourage responses with higher estimated correctness likelihood to receive higher verbalized confidence. Experiments on
reasoning and knowledge-intensive
benchmarks show that our method improves calibration and failure prediction performance while largely preserving answer accuracy. These results demonstrate that verbalized confidence can be more reliably aligned by decoupling confidence estimation from answer generation and optimizing the relative ordering of confidence across responses.
Read the full abstract ▾
All authors Chen Li, Xiaoling Hu, Songzhu Zheng, Jiawei Zhou, Chao Chen
Machine Learning
PyroAdapt: Adapting Wildfire Prediction under Spatial Heterogeneity and Temporal Shift
Enyi Jiang and Wu Sun
Prediction of wildfire occurrence is a rare-event problem compounded by spatial heterogeneity and temporal distribution shift, as fire occurrences are vastly outnumbered by non-occurrences, and predictor–fire relationship varies across space and time. Models trained on historical fire data may perform poorly under new conditions and require adaptation to the target distribution before operational use. We propose PyroAdapt, a pretrain–retrieve–rank framework that adapts a pretrained model to target conditions by retrieving historical locations with similar conditions and fine-tuning on the
retrievals through risk ranking. For spatial adaptation, we condition risk on terrain, ecoregion
embeddings, and fire rates, accounting for spatial context in the
retrieval, and learn risk ordering from same-day fire–nonfire
cell pairs. We compare direct ranking, residual pairwise DPO (RDPO), and selective ranking through a unified score-gap formulation that characterizes their gradient allocation. Over California (discretized into 666 0.25x0.25 grid
cells), these objectives raise daily average precision from 21.62% for continued focal fine-tuning to 24.35–24.57%, and Top5% recall from 18.70% to 22.11–22.79%. Under a fixed daily detection budget of 34
cells (5% area), selective ranking captures 344 additional positive
cell–days. For fires in the top 5%/10%/20% of dry matter consumption, selective ranking raises recall by 39.70/28.18/20.50 percentage points, respectively. Furthermore, rolling evaluations over Yosemite show that the gains from ranking persist under temporal distribution shift. Together, these results show that PyroAdapt prioritizes the most fire-prone locations under a daily budget constraint and detects more extreme fire events.
Read the full abstract ▾
Machine Learning
Learning Minimally Rigid Graphs with High Realization Counts
Oleksandr Slyvka, Jan Rubeš, Rodrigo Alves et al. (+1 more)
For minimally rigid graphs, the same edge-length data can admit multiple realizations (up to translations and rotations). Finding graphs with exceptionally many realizations is an extremal problem in
rigidity theory, but exhaustive search quickly becomes infeasible due to the super-exponential growth of the number of candidate graphs and the high cost of realization-count evaluation. We propose a reinforcement-learning approach that constructs minimally rigid graphs via 0- and 1-extensions, also known as Henneberg moves. We optimize realization-count invariants using the Deep Cross-
Entropy Method with a policy parameterized by a Graph Isomorphism Network encoder and a permutation-equivariant extension-level action head. Empirically, our method matches the known optima for planar realization counts and improves the best known bounds for spherical realization counts, yielding new record graphs.
Read the full abstract ▾
All authors Oleksandr Slyvka, Jan Rubeš, Rodrigo Alves, Jan Legerský
Machine Learning · Cross-listed from cs.CL
ORBIT: Preserving Foundational Language Capabilities in GenRetrieval via Origin-Regulated Merging
Neha Verma, Nikhil Mehta, Shao-Chuan Wang et al. (+7 more)
Despite the rapid advancements in large
language model (LLM) development, fine-tuning them for specific tasks often results in the catastrophic forgetting of their general,
language-based
reasoning abilities. This work investigates and addresses this challenge in the context of the Generative
Retrieval (GenRetrieval) task. During GenRetrieval fine-tuning, we find this forgetting occurs rapidly and correlates with the distance between the fine-tuned and original model parameters. Given these observations, we propose ORBIT, a novel approach that actively tracks the distance between fine-tuned and initial model weights, and uses a weight averaging strategy to constrain model drift during GenRetrieval fine-tuning when this inter-model distance exceeds a maximum threshold. Our results show that ORBIT retains substantial text and
retrieval performance by outperforming both common continual learning baselines and related regularization methods that also employ weight averaging.
Read the full abstract ▾
All authors Neha Verma, Nikhil Mehta, Shao-Chuan Wang, Naijing Zhang, Alicia Tsai, Li Wei, Lukasz Heldt, Lichan Hong, Ed Chi, Xinyang Yi
Machine Learning
Aligning Flow Map Policies with Optimal Q-Guidance
Christos Ziakas, Alessandra Russo and Avishek Joey Bose
Generative policies based on expressive model classes, such as
diffusion and flow matching, are well-suited to complex control problems with highly
multimodal action distributions. Their expressivity, however, comes at a significant inference cost: generating each action typically requires simulating many steps of the generative process, compounding latency across sequential decision-making rollouts. We introduce flow map policies, a novel class of generative policies designed for fast action generation by learning to take arbitrary-size jumps including one-step jumps-across the generative dynamics of existing flow-based policies. We instantiate flow map policies for offline-to-online reinforcement learning (RL) and formulate online adaptation as a trust-region optimization problem that improves the critic's Q-value while remaining close to the offline policy. We theoretically derive FLOW MAP Q-GUIDANCE (FMQ), a principled closed-form learning target that is optimal for adapting offline flow map policies under a critic-guided trust-region constraint. We further introduce Q-GUIDED BEAM SEARCH (QGBS), a
stochastic flow-map sampler that combines renoising with beam search to enable iterative inference-time refinement. Across 12 challenging robotic manipulation and locomotion tasks from OGBench and RoboMimic, FMQ achieves state-of-the-art performance in offline-to-online RL, outperforming the previous one-step policy MVP by a relative improvement of 21.3% on the average success rate.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
Stories in Space: In-Context Learning Trajectories in Conceptual Belief Space
Eric Bigelow, Raphaël Sarfati, Daniel Wurgaft et al. (+5 more)
Large
Language Models (LLMs) update their behavior in context, which can be viewed as a form of
Bayesian inference. However, the structure of the latent hypothesis space over which this inference operates remains unclear. In this work, we propose that LLMs assign beliefs over a low-dimensional
geometric space - a conceptual belief space - and that in-context learning corresponds to a trajectory through this space as beliefs are updated over time. Using story understanding as a natural setting for dynamic belief updating, we combine behavioral and representational analyses to study these trajectories. We find that (1) belief updates are well-described as trajectories on low-dimensional, structured
manifolds; (2) this structure is reflected consistently in both model behavior and internal representations and can be decoded with simple linear probes to predict behavior; and (3) interventions on these representations causally steer belief trajectories, with effects that can be predicted from the geometry of the conceptual space. Together, our results provide a
geometric account of belief dynamics in LLMs, grounding
Bayesian interpretations of in-context learning in structured conceptual representations.
Read the full abstract ▾
All authors Eric Bigelow, Raphaël Sarfati, Daniel Wurgaft, Owen Lewis, Thomas McGrath, Jack Merullo, Atticus Geiger, Ekdeep Singh Lubana
Machine Learning
Predicting Decisions of AI Agents from Limited Interaction through Text-Tabular Modeling
Eilam Shapira, Moshe Tennenholtz and Roi Reichart
AI
agents negotiate and transact in natural
language with unfamiliar counterparts: a buyer bot facing an unknown seller, or a procurement assistant negotiating with a supplier. In such interactions, the counterpart's LLM,
prompts, control logic, and rule-based fallbacks are hidden, while each decision can have monetary consequences. We ask whether an
agent can predict an unfamiliar counterpart's next decision from a few interactions. To avoid real-world logging confounds, we study this problem in controlled bargaining and negotiation games, formulating it as target-adaptive text-tabular prediction: each decision point is a table row combining structured game state, offer history, and dialogue, while $K$ previous games of the same target
agent, i.e., the counterpart being modeled, are provided in the
prompt as labeled adaptation examples. Our model is built on a tabular foundation model that represents rows using game-state features and LLM-based text representations, and adds LLM-as-Observer as an additional representation: a small frozen LLM reads the decision-time state and dialogue; its answer is discarded, and its hidden state becomes a decision-oriented feature, making the LLM an encoder rather than a direct few-shot predictor. Training on 13 frontier-LLM
agents and testing on 91 held-out scaffolded
agents, the full model outperforms direct LLM-as-Predictor prompting and game+text features baselines. Within this tabular model, Observer features contribute beyond the other feature schemes: at $K=16$, they improve response-prediction AUC by about 4 points across both tasks and reduce bargaining offer-prediction error by 14%. These results show that formulating counterpart prediction as a target-adaptive text-tabular task enables effective adaptation, and that hidden LLM representations expose decision-relevant signals that direct prompting does not surface.
Read the full abstract ▾
Machine Learning · Cross-listed from stat.ML
Model-based Bootstrap of Controlled Markov Chains
Ziwei Su, Imon Banerjee and Diego Klabjan
We propose and analyze a model-based bootstrap for transition kernels in finite controlled Markov chains (CMCs) with possibly nonstationary or history-dependent control policies, a setting that arises naturally in offline reinforcement learning (RL) when the behavior policy generating the data is unknown. We establish distributional consistency of the bootstrap transition
estimator in both a single long-chain regime and the episodic offline RL regime. The key technical tools are a novel bootstrap law of large numbers (LLN) for the visitation counts and a novel use of the martingale central limit
theorem (CLT) for the bootstrap transition increments. We extend bootstrap distributional consistency to the downstream targets of offline policy evaluation (OPE) and optimal policy recovery (OPR) via the delta method by verifying Hadamard differentiability of the Bellman operators, yielding asymptotically valid confidence intervals for value and $Q$-functions. Experiments on the RiverSwim problem show that the proposed bootstrap confidence intervals (CIs), especially the percentile CIs, outperform the episodic bootstrap and plug-in CLT CIs, and are often close to nominal ($50\%$, $90\%$, $95\%$) coverage, while the baselines are poorly calibrated at small sample sizes and short episode lengths.
Read the full abstract ▾
Machine Learning
OGLS-SD: On-Policy Self-Distillation with Outcome-Guided Logit Steering for LLM Reasoning
Yuxiao Yang, Xiaoyun Wang and Weitong Zhang
We study on-policy self-
distillation (OPSD), where a
language model improves its
reasoning ability by distilling privileged teacher distributions along its own on-policy trajectories. Despite its promise, OPSD can suffer from training instability due to a pattern mismatch between teacher and student responses. Self-reflected teacher responses may introduce reflection-induced biases and response templates that miscalibrate
token-level supervision, ultimately harming the student's
reasoning ability. To mitigate this issue, we propose OGLS-SD, an outcome-guided logit-steering framework that leverages verifiable outcome rewards to calibrate privileged teacher logits. Specifically, OGLS-SD contrasts teacher logits induced by successful and failed on-policy trajectories, constructing an outcome-discriminative steering direction for
token-level guidance. Experiments on mathematical
reasoning benchmarks show that OGLS-SD stabilizes self-
distillation and improves performance over standard OPSD and other variants.
Read the full abstract ▾
Machine Learning
Detecting overfitting in Neural Networks during long-horizon grokking using Random Matrix Theory
Hari K. Prakash and Charles H Martin
Training Neural Networks (NNs) without overfitting is difficult; detecting that overfitting is difficult as well. We present a novel Random
Matrix Theory method that detects the onset of overfitting in deep learning models without access to train or test data. For each model layer, we randomize each weight
matrix element-wise, $\mathbf{W} \to \mathbf{W}^{\mathrm{rand}}$, fit the randomized empirical spectral distribution with a Marchenko-Pastur distribution, and identify large outliers that violate self-averaging. We call these outliers Correlation Traps. During the onset of overfitting, which we call the "anti-grokking" phase in long-horizon grokking, Correlation Traps form and grow in number and scale as test accuracy decreases while train accuracy remains high. Traps may be benign or may harm generalization; we provide an empirical approach to distinguish between them by passing random data through the trained model and evaluating the JS divergence of output logits. Our findings show that anti-grokking is an additional grokking phase with high train accuracy and decreasing test accuracy, structurally distinct from pre-grokking through its Correlation Traps. More broadly, we find that some foundation-scale LLMs exhibit the same Correlation Traps, indicating potentially harmful overfitting.
Read the full abstract ▾
Machine Learning · Cross-listed from astro-ph.EP
Trajectory-Agnostic Asteroid Detection in TESS with Deep Learning
Brian P. Powell, Jorge Martinez-Palomera, Amy Tuson et al. (+3 more)
We present a novel method for extracting moving objects from TESS data using machine learning. Our approach uses two stacked 3D U-Nets with skip connections, which we call a W-Net, to filter background and identify pixels containing moving objects in TESS image time-series data. By augmenting the training data through rotation of the image cubes, our method is robust to differences in speed and direction of asteroids, requiring no assumptions for either parameter range which are typically required in "shift-and-stack" type
algorithms. We also developed a novel method for learned data scaling that we call Adaptive Normalization, which allows the neural network to learn the ideal range and scaling distribution required for optimal data processing. We built a code for creating TESS training data with asteroid masks that served as the foundation of our effort (tess-asteroid-ml), which we publicly released for the benefit of the community. Our method is not limited to TESS, but applicable for implementation in other similar time-domain surveys, making it of particular interest for use with data from upcoming missions such as the Nancy Grace Roman Space
Telescope and NEOSurveyor.
Read the full abstract ▾
All authors Brian P. Powell, Jorge Martinez-Palomera, Amy Tuson, Christina Hedges, Jessie Dotson, Jordan Caraballo-Vega
Machine Learning · Cross-listed from cs.CV
SEMIR: Semantic Minor-Induced Representation Learning on Graphs for Visual Segmentation
Luke James Miller and Yugyung Lee
Segmenting small and sparse structures in large-scale images is fundamentally constrained by voxel-level,
lattice-bound computation and extreme class imbalance – dense, full-resolution inference scales poorly and forces most pipelines to rely on fixed regionization or downsampling, coupling computational cost to image resolution and attenuating boundary evidence precisely where minority structures are most informative. We introduce SEMIR (Semantic Minor-Induced Representation Learning), a representation framework that decouples inference from the native grid by learning a task-adapted,
topology-preserving latent graph representation with exact decoding. SEMIR transforms the underlying grid graph into a compact, boundary-aligned graph minor through parameterized edge contraction, node deletion, and edge deletion, while preserving an exact lifting map from minor predictions to
lattice labels. Minor construction is formalized as a few-shot structure learning problem that replaces hand-tuned preprocessing with a boundary-
alignment objective: minor parameters are learned by maximizing agreement between predicted boundary elements and target-specific semantic edges under a boundary Dice criterion, and the induced minor is annotated with scale- and rotation-robust
geometric and intensity descriptors and supports efficient region-level inference via message passing on a graph neural network (GNN) with relational edge features. We
benchmark SEMIR on three tumor
segmentation datasets – BraTS 2021, KiTS23, and LiTS – where targets exhibit high structural variability and distributional uncertainty. SEMIR yields consistent improvements in minority-structure Dice at practical runtime. More broadly, SEMIR establishes a framework for learning task-adapted,
topology-preserving latent representations with exact decoding for high-resolution structured visual data.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.MA
Events as Triggers for Behavioral Diversity in Multi-Agent Reinforcement Learning
Hannes Büchi, Manon Flageat, Eduardo Sebastián et al. (+1 more)
Effective multi-
agent cooperation requires
agents to adopt diverse behaviors as task conditions evolve-and to do so at the right moment. Yet, current Multi-
Agent Reinforcement Learning (MARL) frameworks that facilitate this diversity are still limited by the fact that they bind fixed behaviors to fixed
agent identities. Consequently, they are ill-equipped for tasks where
agents need to take on different roles at very specific moments in time. We argue that, to define these behavioral transitions, the missing ingredient is $\textbf{events}$. Events are changes in the state of the system that induce qualitative changes in the task. Based on this view, we introduce a framework that decouples
agent identity from behavior, capturing a continuous
manifold from which
agents instantiate their behaviors in response to events. This framework is based on two elements. First, to build an expressive behavior
manifold, we introduce Neural
Manifold Diversity (NMD), a formal distance metric that remains well-defined when behaviors are transient and
agent-agnostic. Second, we use an event-based hypernetwork that generates Low-Rank Adaptation (LoRA) modules over a shared team policy, enabling on-the-fly
agent-policy reconfiguration in response to events. We prove that this construction ensures that diversity does not interfere with reward maximization by design. Empirical results demonstrate that our framework outperforms established baselines across
benchmarks while exhibiting zero-shot generalization, and being the only method that solves tasks requiring sequential behavior reassignment.
Read the full abstract ▾
All authors Hannes Büchi, Manon Flageat, Eduardo Sebastián, Amanda Prorok
Machine Learning · Cross-listed from cs.SD
A Semi-Supervised Framework for Speech Confidence Detection using Whisper
Adam Wynn and Jingyun Wang
Automatic detection of speaker confidence is critical for adaptive computing but remains constrained by limited labelled data and the subjectivity of paralinguistic annotations. This paper proposes a semi-supervised hybrid framework that fuses deep semantic
embeddings from the Whisper encoder with an interpretable acoustic feature vector composed of eGeMAPS descriptors and auxiliary probability estimates of vocal stress and disfluency. To mitigate reliance on scarce ground truth data, we introduce an Uncertainty-Aware Pseudo-Labelling strategy where a model generates labels for unlabelled data, retaining only high-quality samples for training. Experimental results demonstrate that the proposed approach achieves a Macro-F1 score of 0.751, outperforming self-supervised baselines, including WavLM, HuBERT, and Wav2Vec 2.0. The hybrid architecture also surpasses the unimodal Whisper baseline, yielding a 3% improvement in the minority class, confirming that explicit prosodic and auxiliary features provide necessary corrective signals which are otherwise lost in deep semantic representations. Ablation studies further show that a curated set of high confidence pseudo-labels outperforms indiscriminate large scale augmentation, confirming that data quality outweighs quantity for perceived confidence detection.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
Rui Min, Tianyu Pang, Chao Du et al. (+2 more)
Large
language models (LLMs) have demonstrated remarkable capabilities, but they still frequently produce
hallucinations. These
hallucinations are difficult to detect in
reasoning-intensive tasks, where the content appears coherent but contains errors like logical flaws and unreliable intermediate results. While step-level analysis is commonly used to detect internal
hallucinations, it suffers from limited granularity and poor scalability due to its reliance on step
segmentation. To address these limitations, we propose TokenHD, a holistic pipeline for training
token-level
hallucination detectors. Specifically, TokenHD consists of a scalable data engine for synthesizing large-scale
hallucination annotations along with a training recipe featuring an importance-weighted strategy for robust model training. To systematically assess the detection performance, we also provide a rigorous evaluation protocol. Through training within TokenHD, our detector operates directly on free-form text to identify
hallucinations, eliminating the need for predefined step
segmentation or additional text reformatting. Our experiments show that even a small detector (0.6B) achieves substantial performance gains after training, surpassing much larger
reasoning models (e.g., QwQ-32B), and detection performance scales consistently with model size from 0.6B to 8B. Finally, we show that our detector can generalize well across diverse practical scenarios and explore strategies to further enhance its cross-domain generalization capability.
Read the full abstract ▾
All authors Rui Min, Tianyu Pang, Chao Du, Minhao Cheng, Yi R. Fung
Machine Learning
Let the Batch Decide, On- or Off-Policy: Effective Sample Size for RL Post-Training
Rasool Fakoor, Murdock Aubry, Nicholas Stranges et al. (+1 more)
Reinforcement learning (RL) is structurally harder than supervised learning because the policy changes the data distribution it learns from. The resulting fragility is especially visible in large-model training using nominally on-policy RL, where the training and rollout systems can differ in numerical precision, sampling and other implementation details. The common approach handles this difference by constraining the update with a fixed clip range, so the
algorithm becomes more sensitive to its configuration and must be retuned whenever the task, model scale or distribution mismatch changes. We argue that the fragility traces to two linked concerns that this one fixed range has to cover. The first is a trust-region concern: an update should not move the policy too far from its current value. The second is an off-policy concern: data from older or different behavior policies should influence the update only to the extent that the update remains reliable. Neither concern is a constant to set in advance, and the policy ratios of the current batch show how severe each one is. We adopt P3O, a simple yet effective batch-adaptive objective from prior work that replaces fixed clipping with the normalized effective sample size of the per-
token policy ratios, which we take against the rollout engine's log-probabilities. The same statistic caps the score-function weight and sets the strength of an off-policy regularizer, so the update stays close to the on-policy update when the ratios are close to one. When stale or mismatched data concentrate the ratios, the update tightens but keeps a nonzero score-function weight on every
token. Experiments on math
reasoning show that P3O compares favorably with tuned baselines and, under FP8 rollouts on Qwen3-8B-Base, avoids the late collapse of GRPO. P3O adds no new objective hyper-parameters and removes the clip range.
Read the full abstract ▾
All authors Rasool Fakoor, Murdock Aubry, Nicholas Stranges, Alexander J. Smola
Machine Learning
Discrete Flow Matching for Offline-to-Online Reinforcement Learning
Fairoz Nower Khan, Nabuat Zaman Nahim and Peizhong Ju
Many reinforcement learning (RL) tasks have discrete action spaces, but most generative policy methods based on
diffusion and flow matching are designed for continuous control. Meanwhile, generative policies usually rely heavily on offline
datasets and offline-to online RL is itself challenging, as the policy must improve from new interaction without losing useful behavior learned from static data. To address those challenges, we introduce DRIFT, an online fine-tuning method that updates an offline pretrained continuous time Markov chain (CTMC) policy with an advantage-weighted discrete flow matching loss. To preserve useful pretrained knowledge, we add a path-space penalty that regularizes the full CTMC trajectory distribution, rather than only the final action distribution. For large discrete action spaces, we introduce a candidate-set approximation that updates the actor over a small subset of actions sampled from reference-policy rollouts and uniform exploration. Our theoretical analysis shows that focusing on a small set of high probability actions can reduce computation while keeping the approximation error small, making our method practical for large action spaces. We evaluate on macro-action MinAtar games where offline-pretrained DRIFT outperforms its online-only counterpart across all five games and achieves the highest mean return on three and on oracle free Jericho, DRIFT achieves a mean normalized score of 7.3% with a lightweight GRU encoder, outperforming the evaluated pretrained
language-model baselines. Ablations demonstrate the importance of appropriately weighted path-space regularization and reference-guided candidate construction.
Read the full abstract ▾
Machine Learning
Agent-Based Post-Hoc Correction of Agricultural Yield Forecasts
Matthew Beddows, Aiden Durrant and Georgios Leontidis
Accurate crop yield
forecasting in commercial soft fruit production is constrained by the data available in typical commercial farm records, which lack the sensor networks, satellite imagery, and high-resolution meteorological inputs that most state-of-the-art approaches assume. We propose a structured LLM
agent framework that performs post-hoc correction of existing model predictions, encoding agricultural domain knowledge across tools for phase detection, bias learning, and range validation. Evaluated on a proprietary strawberry yield
dataset and a public USDA corn harvest
dataset,
agent refinement of XGBoost reduced MAE by 20% and MASE by 56% on strawberry, with consistent improvements across Moirai2 (MAE 24%, MASE 22%) and Random Forest (MAE 28%, MASE 66%) baselines. Using Llama 3.1 8B as the
agent produced the strongest corrections across all configurations; LLaVA 13B showed inconsistent gains, highlighting sensitivity to the choice of refinement model.
Read the full abstract ▾