Machine Learning
Bellman-Certified Rounding for Sparse Policy Deployment in MDPs
Zhaojun Peng
Continuous policy optimization may spread an update across many states, even when deployment permits only a few complete state-level changes. We study how much discounted return can be retained when continuous row mixtures are rounded to sparse binary policies in finite MDPs. Policy-dependent visitation couples the row edits, while long horizons make global
curvature bounds conservative. From $2d+2$ Bellman solves, we derive reusable envelopes that support uniform and candidate-specific guarantees before rounding. A rank-two rational representation of each exchange further permits weighted
curvature integration along the realized trajectory. We prove that linear dimension dependence is unavoidable when the budget scales, and that exact global
curvature thresholding is hard. Candidate-specific bounds raise pre-rounding certification coverage from $48.2\%$ to $74.1\%$ on the structured suite. At $\gamma=0.95$, local integration lowers the median bound-to-loss ratio from $402.3$ to $2.08$ on coupled instances.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
CAST: Cost-Aware Speculative Trees from One-Pass Block Drafters
Jungseob Lee and Sugyeong Eo
Speculative decoding accelerates large
language model inference by drafting future
tokens cheaply and verifying them with the target model in parallel. Block drafters score a whole block of future
tokens in one forward pass, yet standard decoding verifies only the top-scoring chain and discards the other candidates. Because these candidates are already scored, verifying more of them adds target computation but no extra drafting. We introduce CAST (Cost-Aware Speculative Trees), which packs these candidates into a tree and verifies it in a single target pass, leaving the target model, drafter weights, and decoding rule untouched. To decide how wide the tree should be, CAST adds candidates while the expected gain from the next one outweighs the verification time it adds. The width therefore adapts to each deployment from a latency measurement, without sweeping over widths. We evaluate CAST across five domains on three GPU generations and two model families. At its predicted width, CAST is faster than the standard chain in all eight settings, by up to 43%. We also find that the best width depends strongly on the deployment. Where verification cost jumps at a kernel boundary, a 128-
token tree is only 2% faster than the standard chain, whereas the tree at the predicted width is 20% faster. Furthermore, we prove that CAST leaves the target output distribution unchanged under both greedy and sampled decoding. Code is available at https://github.com/js-lee-AI/CAST.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.CL
Refusal Localizes, the Damage Relocates: Safety Layers Under Few-Sample Fine-Tuning
Jungseob Lee, Dongyub Jude Lee, Sugyeong Eo et al. (+3 more)
Fine-tuning adapts aligned large
language models (LLMs) to downstream tasks, but a few dozen harmful examples can remove their refusal of harmful requests. Prior work localizes safety-related behavior to specific layers, directions, and
tokens, suggesting targets for protection. We test whether successful localization and recovery support defenses that survive changes in the attack. Across six checkpoints from four model families, harmful and benign
prompts remain linearly separable after attack, and patching full clean hidden states into the compromised model restores refusal at a reproducible transition depth. Building on a prior layer-freezing defense, we freeze every layer up to this depth and repeat the attack. At a hundred harmful examples, refusal remains near zero on all six checkpoints, with recovery transitions above the frozen boundary. In a second study, removing the update's top two singular directions restores refusal after short
attention-only fine-tunes on four checkpoints. On Llama-3.1-8B, ordinary training changes weaken this repair and an attacker who spreads the update defeats it. A spectral detector calibrated on benign Llama fine-tunes misses most repair failures on that checkpoint. Localized freezing can nevertheless help preserve refusal when a few harmful examples enter training data unintentionally. These results show that an attacker can bypass a region identified by recovery and defeat a repair that works across multiple checkpoints, motivating five checks for defenses against adaptive fine-tuning. Code is available at https://github.com/js-lee-AI/refusal-relocates.
Read the full abstract ▾
All authors Jungseob Lee, Dongyub Jude Lee, Sugyeong Eo, Seongtae Hong, Seungyoon Lee, Heuiseok Lim
Machine Learning · Cross-listed from cs.RO
DriftOPD: Sequence-Level Reverse-KL Distillation for One-Step VLA Policies
Youngjun Jun, Kyumin Choi, Youngmin Kim et al. (+4 more)
Vision-
Language-Action (VLA) models increasingly rely on action experts that generate short action chunks under receding-horizon control. While chunk-level training is convenient across
robot embodiments, it optimizes local action likelihood without explicitly accounting for long-horizon task success. Sequence-level reinforcement learning can address this limitation, but typically requires policy rollouts and closed-loop interaction, which are costly for real-
robot manipulation. We introduce DriftOPD, a teacher-free, rollout-free framework for sequence-level on-policy
distillation of continuous VLA action experts. We show that the sequence-level reverse Kullback-Leibler (KL) divergence decomposes into a chunk-level reverse-KL term and a future-potential term that captures the long-horizon effect of the current action. DriftOPD optimizes these two terms using a one-step drifting objective and a Q-function critic learned from offline demonstrations, respectively, enabling sequence-level optimization with only offline data and one-step action generation. Across multiple VLA architectures in simulation and real-world manipulation, DriftOPD generally outperforms existing one-step
distillation baselines while achieving task success performance comparable to multi-step teacher policies. These results demonstrate that long-horizon behavior can be effectively distilled into one-step VLA action experts without online interaction or a separate teacher.
Read the full abstract ▾
All authors Youngjun Jun, Kyumin Choi, Youngmin Kim, Seonghyun Jin, Sunwoo Park, Jangho Park, Jong Chul Ye
Machine Learning · Cross-listed from cs.CV
Beyond Pixel Reconstruction: Retrieval-Guided Glyph-Aware Restoration for Low-Resource Manchu Historical Documents
Ting Huang, Dongdong Wang, Mingqiu Liang et al. (+1 more)
Historical Manchu documents preserve invaluable linguistic and cultural heritage, yet their digitization is hindered by severe degradations and the scarcity of paired training data. Existing document restoration methods primarily optimize pixel-level reconstruction, which can produce visually plausible results while failing to preserve the structural identity of Manchu glyphs. To address this limitation, we propose a
retrieval-guided glyph-aware restoration framework that goes beyond pixel reconstruction by explicitly incorporating glyph-level structural knowledge. Our method retrieves relevant glyph exemplars to provide structural guidance during restoration and integrates this information into the reconstruction process, improving the recovery of degraded character structures under low-resource conditions. Extensive experiments on Manchu historical documents demonstrate that the proposed approach improves both image restoration quality and glyph-level fidelity compared with existing restoration methods. These results highlight the importance of incorporating character-aware structural priors for reliable restoration of low-resource historical documents.
Read the full abstract ▾
All authors Ting Huang, Dongdong Wang, Mingqiu Liang, Siyang Lu
Machine Learning
In-Distribution Imagination for Model-Based Offline Reinforcement Learning
Mintae Kim and Koushil Sreenath
Model-based offline reinforcement learning (MBORL) improves sample efficiency through model-generated trajectories. However, accumulative model error can drive imagined trajectories outside the offline data distribution, leading to unrealistic synthetic data and unstable policy optimization. Many existing methods primarily control rollouts using transition-level uncertainty. We propose in-distribution imagination (IDI), a rollout control framework that estimates trajectory support in a learned representation space and adaptively truncates rollouts that leave the offline trajectory
manifold. Combined with trajectory-regularized RL, an extension of
entropy-regularized RL, IDI consistently improves performance in limited-data settings. Experiments show that trajectory support predicts rollout failure substantially better than transition-level uncertainty, highlighting the importance of trajectory-level rollout control in MBORL.
Read the full abstract ▾
Machine Learning
Provable Test-Time Scaling for Beam Search in LLM Reasoning
Qijia He, Yu Huang, Yuan Cheng et al. (+2 more)
Beam-search-based test-time methods provide an effective way to improve large
language model (LLM) performance on long-horizon generation by
pruning invalid
reasoning paths early, leading to significantly improved
reasoning efficiency and more favorable test-time cost scaling. Despite strong empirical success, the theoretical understanding of beam search remains limited. In this paper, we study the test-time compute guarantee of the commonly used beam search framework that uses the model's internal log-likelihood for intermediate scoring, while relying on an external reward model only after a complete response is generated. We first establish a lower bound for vanilla beam search, showing that at least $\Omega(C^\star(x)^2)$ samples are required for the optimal response to survive, where $C^\star(x)$ is the
token-level coverage coefficient for
prompt $x$. This motivates our modified confidence-filtered beam search (CF-Beam), which reduces the sufficient coverage dependence from quadratic to nearly linear under prefix competitiveness, for fixed horizon, gap, and target accuracy. We then show that the regret of CF-Beam is upper-bounded by the probability of rare failure events and the reward estimation error scaled by a path-level coverage coefficient, where the rare-failure term vanishes as per-step sampling increases. Our results highlight a fundamental advantage of beam search over sequence-level inference methods such as Best-of-N and Best-of-Majority. While the guarantees of these approaches typically involve coverage coefficients that grow exponentially with the horizon $L$, CF-Beam controls the dominant search-induced term through a
token-level coverage coefficient that scales polynomially with $L$. Our numerical experiments further confirm that beam search is more robust on hard instances and under increasing
reasoning horizons.
Read the full abstract ▾
All authors Qijia He, Yu Huang, Yuan Cheng, Yuxin Chen, Yingbin Liang
Machine Learning
Understanding Off- vs On-Policy Distillation: A Tale of Distinct Training Objectives
Qiwei Di, Xuheng Li, Kaixuan Ji et al. (+3 more)
On-policy
distillation (OPD) learns from teacher feedback on student-generated responses and has shown promise in reducing forgetting relative to supervised fine-tuning (SFT). However, its benefits and fragility remain incompletely understood. We study sequential
distillation from multiple teachers, where the student minimizes its average divergence from the teachers. Forward Kullback–Leibler (KL) divergence yields a weighted arithmetic mixture, while reverse KL yields a normalized weighted
geometric aggregate. We develop
algorithms that learn these targets under off-policy and on-policy feedback, respectively, establishing logarithmic regret bounds in the tabular setting and extending the analysis to function approximation. By analyzing these aggregation targets, we identify mechanisms that help explain both the benefits and fragility of OPD. Relative to forward KL, reverse KL can better retain a confident expert's preferences under uninformative feedback, but is more sensitive to teachers that assign very low probabilities to correct responses. Its
token-level conditionals also reveal a dependence on continuation distributions that can favor incorrect prefixes over long horizons.
Read the full abstract ▾
All authors Qiwei Di, Xuheng Li, Kaixuan Ji, Chenggong Zhang, Heyang Zhao, Quanquan Gu
Machine Learning · Cross-listed from stat.ML
Bandits with Multiple Optimal Arms: Minimax Regret and Non-Adaptivity
Kaixuan Ji, Qiwei Di, Qingyue Zhao et al. (+2 more)
We study multi-armed bandits (MAB) with multiple optimal arms, motivated by the fact that many practical decision making problems admit multiple correct answers. For $K$-armed bandits with $A$ optimal arms, we first provide a sharper analysis of previous sub-sampling
algorithms (De Heide et al., 2021; Zhu and Nowak, 2020), establishing a $\tilde{O}\Big(\frac{K-A}{\sqrt{KA}}\sqrt{T} \Big)$ minimax regret, where $T$ is the total number of interactions and $\tilde O(\cdot)$ drops all constant and logarithmic factors, improving the previous $\tilde{O}(\sqrt{KT/A})$ regret. We then provide a matching lower bound up to logarithmic factors, indicating that our established rate is nearly minimax-optimal. We further show that the knowledge of $A$ up to $\tilde{O}(1)$ factors is necessary to achieve near-optimal regret, as near-optimal
algorithms for one number of optimal arms must incur substantially larger regret than optimal regret for a smaller number. Overall, our results provide a comprehensive minimax characterization of $K$-armed bandits with $A$ over the entire range of $1 \leq A \leq K-1$.
Read the full abstract ▾
All authors Kaixuan Ji, Qiwei Di, Qingyue Zhao, Heyang Zhao, Quanquan Gu
Machine Learning · Cross-listed from cs.RO
TERRA: Terrain-Aware Reconstruction, Retargeting and Control for Musculoskeletal Locomotion
Merkourios Simos, Chengkun Li, Bianca Ziliotto et al. (+1 more)
Recent advances in musculoskeletal modeling and reinforcement learning have enabled muscle-actuated
agents to reproduce increasingly complex human motions. Yet these capabilities remain largely confined to flat ground, in part because motion
datasets rarely include aligned terrain geometry and because retargeting terrain interactions to complex musculoskeletal bodies is challenging. We present TERRA, an end-to-end pipeline for terrain-aware retargeting and control of musculoskeletal locomotion. From kinematic trajectories alone, TERRA combines terrain priors, estimated contacts, and negative free-space evidence to recover task-relevant support geometry. TERRA further considers anatomical, tendon-continuity, and contact constraints during retargeting. Using the resulting motion-terrain pairs from five
datasets, we successfully train a single muscle-actuated control policy on 9.4 hours of diverse locomotion. Across reconstruction, retargeting, and held-out tracking
benchmarks, TERRA improves terrain accuracy, sharply reduces anatomical and interaction violations, and achieves the highest observed completion rate over supported terrain families. Overall, TERRA provides a practical route from scene-less motion data to muscle-actuated locomotion over diverse non-flat terrain. Project website: https://cnai.epfl.ch/terra/
Read the full abstract ▾
All authors Merkourios Simos, Chengkun Li, Bianca Ziliotto, Alexander Mathis
Machine Learning
Uncertainty-Normalized Margins for Direct Preference Optimization
Sadegh Khorasani, Petrus Mikkola and Matthias Grossglauser
Direct preference optimization (DPO) models binary preferences through a Bradley-Terry model with a common noise scale, without explicitly accounting for preference strength or
prompt-dependent uncertainty from human feedback. We introduce uncertainty-normalized margin DPO (UNM-DPO), which combines strength-dependent margins with a learned
prompt scale. Motivated by a heteroskedastic Bradley-Terry model, we develop two training objectives. Both compare the implicit rewards of preferred and rejected responses, derived from response log-probability ratios to a reference policy. Advantage-only (AO) divides this reward difference by the
prompt scale before subtracting the margin; whole-residual (WR) subtracts the margin before dividing by the scale. For the WR comparison model, we establish a necessary and sufficient condition under which known margins make the
prompt scale identifiable. We introduce a practical procedure for learning the scale. Building on WR, we introduce ULNM-DPO-WR, which normalizes each response's implicit reward by its length. We evaluate our methods against DPO and related baselines on HelpSteer2 and HelpSteer3, using the Skywork reward model as a judge. With Llama-3.1-8B-Instruct, ULNM-DPO-WR achieves tie-adjusted win rates against matched DPO of 68.00% and 65.31% on evaluation panels, with higher mean rewards and shorter responses on average. On AlpacaEval with a GPT-4.1 judge and GPT-4-Turbo reference answers, the same 8B policy achieves a length-controlled win rate of 21.62%, compared with 16.39% for DPO and 15.30% for SimPO. These results demonstrate the potential of combining preference-strength margins, learned
prompt scales, and length normalization for policy optimization.
Read the full abstract ▾
Machine Learning
Alignment via Training Against Probes Without Losing Monitorability
Lena Libon, Alexander Panfilov, Ben Rank et al. (+3 more)
Models are usually aligned based on their observed outputs, using demonstrations, preference data, or reward signals. These objectives reward responses that look aligned. More capable models may learn to satisfy them without internalizing the intended behavior, for example by faking compliance during training. Such superficial compliance could be harder when the objective is defined on model internals rather than outputs. Therefore, we study probe-guided fine-tuning, using probes that detect undesired properties in model activations as a direct training signal. We evaluate linear and non-linear probes with different numbers of probes per layer across two
alignment objectives: harmlessness and honesty. We find that training against probes that do not update during training is an easily exploitable objective, while continuously updated probes substantially reduce harmfulness and improve honesty while preserving utility. Probe-guided fine-tuning achieves better safety-utility trade-offs than DPO and inference-time steering, while being substantially more robust against jailbreak and abliteration attacks. Moreover, the concepts stay linearly encoded after fine-tuning, meaning oversight is not lost by our method. Training against probes thus offers a way to shape what models represent rather than only what they output, which may become increasingly important as models get better at making their outputs look aligned.
Read the full abstract ▾
All authors Lena Libon, Alexander Panfilov, Ben Rank, Xin Chen, Jonas Geiping, Maksym Andriushchenko
Machine Learning · Cross-listed from cs.CE
SHIFT-Truck: A High-Fidelity Aerodynamics Dataset and Benchmark for Pickup Trucks
Riddhiman Raut, Yin Yu, Aashwin Anand Mishra et al. (+4 more)
Pickup trucks account for 14% of new light-duty vehicles produced in the United States, yet are among the least aerodynamic. Their open cargo bed adds a flow absent from existing automotive aerodynamics
datasets such as DrivAerML and SHIFT-SUV: the shear layer leaving the cab roof passes over a recirculating bed flow before separating again at the tailgate. The resulting drag lowers fuel efficiency, raises emissions and limits the range of electric trucks. Scale-resolved Computational Fluid Dynamics (CFD) is too costly for broad design exploration; neural surrogates can predict flow features at a fraction of that cost, provided they are trained on large-scale, high-fidelity, domain-specific data. We introduce SHIFT-Truck, the first such
dataset for pickup trucks. It comprises 1,000 Spalart-Allmaras delayed detached-eddy simulations (SA-DDES) of a reference pickup geometry morphed across 17 shape parameters. Each case is run on a mesh of about 100 million
cells at a Reynolds number of $1.4 \times 10^7$ and released with time-averaged surface pressure, wall shear stress, volumetric pressure and velocity. The setup is verified by grid refinement and repeated runs, and checked against wind-tunnel measurements. We define geometry-grouped splits and
benchmark four neural surrogates, DoMINO, GeoTransolver, AB-UPT and SMART, on surface and volume tracks. SHIFT-Truck also introduces controlled distribution shifts in the operating point, the input surface discretization and the vehicle archetype. Models with strong in-distribution performance can degrade substantially under these shifts: operating-condition changes expose failures to infer speed dependence, while tessellation and cross-vehicle shifts reveal markedly different robustness across architectures. SHIFT-Truck is thus a
benchmark not only for surrogate accuracy but also for generalization across physical and numerical distributions.
Read the full abstract ▾
All authors Riddhiman Raut, Yin Yu, Aashwin Anand Mishra, Michael Emory, Thomas Economon, Peter Lyu, Juan J. Alonso
Machine Learning · Cross-listed from cs.CV
STEPS: Scene Text Editing with Preserved Style Using Diffusion and Contrastive Style Encoding
Nicolas Thiebaut, Nameer Hirschkind, Xiao Yu et al. (+1 more)
We introduce Scene Text Editing with Preserved Style (STEPS), a novel
diffusion model architecture for quality text replacement in images. Scene Text Editing (STE), also known as Visual Text Editing, consists of changing the textual content in an image while conserving the original style, e.g. font, colors, orientation, background, etc. STEPS advances the state of the art in STE through directed focus on improved style preservation. We introduce a style encoder for visual text that captures style independently of textual content, and a model architecture that combines the style encoder with multiple semantic conditions (target text characters encoding and rendered glyphs). STEPS achieves superior results to previous STE methods in style preservation, output readability, and subjective quality.
Read the full abstract ▾
All authors Nicolas Thiebaut, Nameer Hirschkind, Xiao Yu, Kyle Spence
Machine Learning
Proper Scoring Rule-based Diffusion for Probabilistic Weather Forecasting
Joonhyeong Park, Giung Nam, Hyungi Lee et al. (+3 more)
Recent probabilistic weather forecasters train
stochastic predictors with the continuous ranked probability score (CRPS) to generate each ensemble member in a single forward pass. These models learn the predictive distribution from the forecast context alone, which becomes difficult at longer forecast horizons where uncertainty is high. To learn the predictive distribution more effectively, we introduce auxiliary conditional denoising tasks that predict the same future state from the context and its corrupted version, which provides partial future information that can reduce prediction ambiguity. Building on distributional
diffusion models, we learn the conditional distributions of these tasks with a single
stochastic predictor by minimizing a proper scoring rule across noise levels. At inference, the predictor can still generate each ensemble member in a single forward pass at the fully corrupted endpoint. Standard CRPS training is recovered as the endpoint-only special case of our formulation, so our framework extends existing CRPS-based forecasters with only additional conditioning inputs. Controlled experiments show that the auxiliary tasks improve one-step forecasting across architectures, with larger gains at longer forecast horizons. The gains extend to high-dimensional global weather forecasting under both training from scratch and fine-tuning, along with improved calibration and potential benefits for generalization under distribution shift.
Read the full abstract ▾
All authors Joonhyeong Park, Giung Nam, Hyungi Lee, Kyunghyun Cho, Byoungwoo Park, Juho Lee
Machine Learning
Interpretable but Fragile? Robustness of Concept Bottlenecks under Geometric-Semantic Perturbations
Hanwei Zhang, Tianma Hu, Gaojie Jin et al. (+2 more)
Concept Bottleneck Models (CBMs) are designed to provide interpretable intermediate representations, yet how such bottlenecks affect robustness remains unclear, with existing studies reporting mixed and sometimes contradictory findings. We argue that these discrepancies arise from conflating different robustness notions and perturbation regimes, rather than from fundamental disagreements about CBMs themselves. To disentangle these factors, we introduce a generator-based evaluation framework that enables controlled comparisons between standard classifiers and CBMs under two distinct perturbation types: continuous
geometric perturbations in latent space and discrete semantic interventions in concept space. Within this framework, we evaluate robustness both empirically, via prediction and concept-level sensitivity metrics, and certifiably, using randomized smoothing in latent and concept spaces. Across experiments, we reconcile previously conflicting findings by clarifying when, and in what sense, concept bottlenecks do or do not improve robustness. By further analyzing robustness under varying task conditions, including class semantic similarity and concept vocabulary size, we show that interpretability does not inherently confer robustness. Instead, concept bottlenecks shift where and how sensitivity manifests, revealing a nuanced interpretability robustness trade off that depends critically on the perturbation regime and task structure. Together, our results show that interpretability and robustness are distinct objectives: interpretable intermediate representations do not uniformly improve robustness, but instead redistribute sensitivity across perturbation spaces and model families.
Read the full abstract ▾
All authors Hanwei Zhang, Tianma Hu, Gaojie Jin, Xu Cheng, Ronghui Mu
Machine Learning
Geometry-physics confounding impairs PDE learning across varying domains
Yinghao Cheng, Gengxiang Chen, Xu Liu et al. (+6 more)
Learning partial differential equation (PDE) dynamics across varying domains is central to predictive modelling and data-driven discovery of governing equations. However,
geometric variation alters both field representation and the governing differential operators, confounding
geometric effects with intrinsic physical properties in the observed dynamics. This work identifies geometry-physics confounding as a unified failure mechanism for PDE learning across varying domains. In forward operator learning, this confounding increases the burden of inferring geometry-dependent operator changes from finite data, reducing data efficiency and generalisation. In equation discovery, omitting geometry-induced operators misspecifies the candidate library, leading to biased parameters, missed governing terms and spurious terms. We propose a de-confounding framework that makes the known geometry-to-operator transformation explicit. Geometry-induced coefficient fields improve prediction and data efficiency across five operator-learning
benchmarks, while geometry-complete candidate libraries recover the generating equations and reduce held-out PDE residuals by more than two orders of magnitude in both evolving-domain systems. By separating known
geometric action from intrinsic physics, the proposed framework supports more reliable and data-efficient PDE learning across scientific and engineering problems with varying geometries.
Read the full abstract ▾
All authors Yinghao Cheng, Gengxiang Chen, Xu Liu, Qinglu Meng, Yixin Jing, Xiangguo Tang, Wenping Mou, Lihui Wang, Yingguang Li
Machine Learning
Differentiable Structure Learning for Cyclic Linear Gaussian Models with Latent Confounders
Sadegh Khorasani, Ali Najar, Saber Salehkaleybar et al. (+1 more)
We study causal structure learning from observational data in linear Gaussian structural causal models in the presence of directed cycles and an unknown number of exogenous latent confounders, bounded by a given maximum. We derive the covariance of the observed variables and introduce marginal quasi-equivalence, which characterizes when different causal models share a full-dimensional subset of the observational distributions they can generate. We formulate structure learning as minimization of the Gaussian negative log-likelihood with a logarithmically scaled complexity penalty that counts directed edges and latent variables. For a fixed number of observed variables and a fixed upper bound on latent variables, we establish consistency of global score minimizers up to marginal quasi-equivalence under algebraic faithfulness, structural minimality, and model-overlap assumptions. We parameterize the inclusion of directed edges and candidate latent variables using Bernoulli gates, whose continuous probabilities are optimized jointly with the structural coefficients. Averaging the penalized negative log-likelihood over these gates yields an objective with a closed-form differentiable complexity penalty. We prove that this expected objective has the same global infimum as the corresponding discrete structure-learning objective. Experimental results show that our approach achieves lower recovery error than previous methods in several experimental settings.
Read the full abstract ▾
All authors Sadegh Khorasani, Ali Najar, Saber Salehkaleybar, Negar Kiyavash
Machine Learning · Cross-listed from cs.RO
Correcting WHERE, Preserving HOW: Compositional Generalization for Vision-Language-Action Models via Referential Guidance
Yanyan Zhang, Disheng Liu, Xinpeng Li et al. (+8 more)
While
Vision-
Language-Action (VLA) models enable flexible action generation, their generalization across diverse environmental elements, including manipulated objects, destinations, and backgrounds, is limited by the lack of diversity in robotic training data. Trained end-to-end on such data, VLAs tend to exploit visual shortcuts, associating actions with task-irrelevant visual features rather than the intended task semantics. These shortcuts block recomposition of elements already seen by the policy, that is, compositional generalization. Existing approaches mitigate such
entanglement through task-relevant perception or targeted data diversification, but offer no explicit mechanism for unseen recomposition and require backbone-specific modifications with retraining. We observe that under such recomposition, VLAs often fail at global grounding while retaining local manipulation skills that recover near the correct target in familiar configurations. Therefore, we propose Referential Guidance (ReGuide), a training-free wrapper that, given object poses from a grounding module, combines semantic and
geometric rebinding to guide the end-effector into demonstration-supported configurations of the instructed referent, where the frozen policy can resume execution. Experiments in simulation across multiple VLA backbones as well as on a real
robot show that ReGuide improves success rates under compositional shifts by up to 56.8 and 75.0 percentage points, respectively, while preserving standard-task performance.
Read the full abstract ▾
All authors Yanyan Zhang, Disheng Liu, Xinpeng Li, Chaoda Song, Mohsen Hariri, Debargha Ganguly, Wang Yang, Kai Ye, Bryce Grant, Vipin Chaudhary, Yu Yin
Machine Learning · Cross-listed from physics.ao-ph
Methodological Changes to the Attention ResUNet Hourly Precipitation Postprocessor
Thomas M. Hamill
This note is a technical companion to a previously published preprint describing an
Attention Residual U-Net that postprocesses deterministic forecasts from The Weather Company's Global and Regional Atmospheric Forecast (GRAF) model into probabilistic hourly precipitation forecasts. It documents what has changed in that method since publication. Feature-wise Linear Modulation conditioning on calendar season and forecast lead time is used to produce a single trained model for each season, replacing 192 separately trained per-month, per-lead checkpoints. Lead time is extended from 48 to 72 h. Two new input channels are used, per-pixel local solar hour and a static, monthly-varying precipitation climatology. During verification, the climatological reference against which the Brier Skill Score is computed now has an added diurnal dimension, on top of the monthly resolution it already had. Brier Skill Score and reliability are compared between the new vs. the previous training. Forecasts generated with the new training show a modest, consistent improvement of the current training over the original.
Read the full abstract ▾
Machine Learning
Learning-Enabled Estimation: Tight Characterizations under Sample Selection Biases
Vikram Kher, Jane H. Lee, Anay Mehrotra et al. (+1 more)
When can we learn from biased samples? We study regression when outcomes are observed only after passing through selection filters that depend on both
covariates and outcomes themselves, a ubiquitous challenge spanning clinical trials with patient dropout, labor
markets with self-selection, and auctions with strategic entry. Ignoring such selection yields systematically biased conclusions with real-world consequences. This challenge has a long history in econometrics and statistics, starting with Heckman's seminal two-stage model and followed by numerous generalizations. While these works provide various sufficient conditions for identification, a complete characterization of when such regression is possible has remained elusive. In this work, we provide a characterization for when regression is possible in the presence of sample selection bias. Our results establish the minimal assumptions required on the functional forms of selection processes under which regression remains possible, which are particularly relevant in modern settings where selection mechanisms are increasingly complex and opaque. As a corollary of our characterization, we show that there are settings where the regression function can be identified even when the selection filter itself cannot. This observation already goes beyond the “estimate selection filter, then debias regression” paradigm that is followed by virtually all existing approaches. Under natural strengthenings of our identification conditions, we also establish finite-sample estimation guarantees with explicit convergence rates and provide oracle-efficient
algorithms. This yields the first general-purpose estimation method for this broad class of selection problems. Finally, we explore the implications of our results for several well-studied econometric settings with complex selection mechanisms such as auctions with entry costs and labor
markets.
Read the full abstract ▾
All authors Vikram Kher, Jane H. Lee, Anay Mehrotra, Manolis Zampetakis
Machine Learning · Cross-listed from cs.CL
Beyond Oracle Communication: Benchmarking Interactive Intent Alignment Under Miscommunication and Evolving User Intent
Zheyuan Zhang, Mengyuan Chao, Ke Xiao et al. (+5 more)
Modern LLM
agents increasingly tackle complex tasks through interactive, long-horizon exchanges with users, while existing
benchmarks generally assume that users always accurately and sufficiently communicate a fixed intent. However, this oracle communication assumption rarely holds in practice: users may miscommunicate, change their goals, and run out of patience. We define this task setting as Interactive Intent
Alignment, where
agents must recover and continuously track the user's current intent despite imperfect communication and evolving goals. To study this setting, we introduce Drift-Bench++, a principled
benchmark construction pipeline for verified executable tasks with controlled misalignment and intent shifts, along with an interaction protocol featuring finite patience, diverse simulated users, and silent interaction-conditioned shifts. We further develop GRIP, a comprehensive evaluation protocol covering task grounding, user realism, inquiry effectiveness, and adaptation to evolving intent. Across diverse environments, models, and interaction conditions, stronger interaction consistently helps but remains far from oracle performance; Validation on deployed ProdAgent sessions further shows that the modeled failures are prevalent and consequential in deployment. By providing a unified, executable
benchmark for interactive intent
alignment, Drift-Bench++ offers a foundation for evaluating and advancing
agents under realistic communication and evolving intent.
Read the full abstract ▾
All authors Zheyuan Zhang, Mengyuan Chao, Ke Xiao, Ziyi Chen, Daoan Zhang, Yan Zhang, Yanfang Ye, Wei Xu
Machine Learning
JARQ: Joint Alternating Refinement for Quantization
Xinyu Wang, Sicheng Lyu and Xiao-Wen Chang
Group-wise post-training quantizers for large
language models round weights onto a grid that is not refit to the resulting integer codes. We show that this leaves accuracy on the table: the best grid depends on the codes, input correlations couple the errors of different groups, and useful code changes often involve many codes at once. We propose JARQ , a plug-in refinement that starts from any group-wise quantizer and alternates a joint least-squares fit of all group scales with bounded Babai proposals that move many codes of a group together on the current grid. The problem is a bilinear box-constrained mixed-integer least-squares problem; the solver is backpropagation-free, does not increase the layer-wise objective under exact scale solves, and keeps the host's bit width, groups, zero points, and inference cost. Across Llama-2, Llama-3, and Qwen models with RTN, GPTQ, OmniQuant, and AWQ hosts, JARQ lowers perplexity in 90 of 96 comparisons, cuts three-bit RTN perplexity by up to 36%, raises mean multiple-choice accuracy in 23 of 24 configurations, and improves QEP, QuaRot, and OJBKQ outputs, at under a minute per 7B block.
Read the full abstract ▾
Machine Learning
Reinforcement Learning with Complex (valued) Memories
Sathya Kamesh Bhethanabhotla, Efstratios Gavves and André Biedenkapp
Partially observable environments pose a fundamental challenge in deep reinforcement learning, requiring
agents to compress temporal information from observations and maintain a
memory to make effective decisions. While there exist many approaches ranging from gated recurrence to
attention mechanisms and model-based RL, the search for effective representational techniques that can capture long-term dependencies remains an active area of research. In this work we revisit
Unitary recurrent networks (uRNNs) [Arjovsky et al., 2016, Jing et al., 2017], that demonstrated superior gradient flow and associative recall, expressing the recurrence and the hidden state in a complex vector space. Their norm preserving
unitary dynamics enable information propagation through long sequences. To this end, we propose three different versions of uRNNs as drop-in replacements for recurrent PPO architectures, and demonstrate that the simple recurrence and the added degree of freedom from the phase of the complex representations enable significant gains over baselines on several
memory-improvable tasks, including continuous control. We further explore how to preserve the phase information of the complex hidden state for a phase-aware policy by drawing a parallel to how
quantum states are measured. With our methods reaching up to 2-3 $\times$ the reward in environments like rocksample and Craftax compared to the baselines, this work points towards an exciting new direction of representations for RL and the problem of partial observability. Code is available at: https://github.com/Sathya98/qurl
Read the full abstract ▾
Machine Learning
NeurDuo-EEG: A Long-Sequence EEG Foundation Model with Persistent State and Explicit Memory
Yifan Wang, Haiping Liu, Yang Cui et al. (+13 more)
Electroencephalography (
EEG) is recorded continuously over hours, with relevant dynamics spanning timescales from milliseconds to hours. Most
EEG foundation models nevertheless process fixed windows independently, limiting their ability to capture information encoded in long-timescale dynamics. State-space architectures enable persistent recurrent processing, but long-range information remains implicitly compressed in recurrent states. We present NeurDuo-
EEG, a causal
EEG foundation model with channel-resolved persistent
memory. NeurDuo-
EEG introduces multi-timescale
memory management with learned consolidation and selective
retrieval, enabling persistent modelling of continuous
EEG with fixed-size state. It is pre-trained on 3,955 hours of
EEG from 17 public
datasets using multichannel autoregressive prediction of discrete spectral codes. Across three short-window and two long-sequence downstream tasks, NeurDuo-
EEG achieves the best performance on four of five
benchmarks, including all three short-window tasks and seizure detection, where AUC-PR improves from $0.285$ to $0.471$ over the strongest non-NeurDuo baseline. NeurDuo-
EEG also remains competitive on sleep staging and supports efficient streaming inference, with nearly constant per-chunk latency as the available history grows to one hour. Notably, the Small variant achieves this with only 4.7M backbone parameters. These results demonstrate the value of persistent, multi-timescale modelling for both long-sequence and short-window
EEG analysis. Our code is available at https://github.com/YifaNNW/NeurDuo-
EEG.
Read the full abstract ▾
All authors Yifan Wang, Haiping Liu, Yang Cui, Wenhao Cai, Shuhang Li, Xiaoyang Huang, Xianyang Liu, Jingyu Sun, Yizheng Sun, Cunhang Fan, Tianming Du, Jiancheng Yang, Zhenhong Li, Yunhao Zhang, Hongpeng Zhou, Jingyuan Sun
Machine Learning
FlexRouter: Learning Complementary Model Sets for Flexible LLM Routing
Wang Wei, Harry Yang, Tiankai Yang et al. (+6 more)
Existing Large
Language Model (LLM) routing methods score LLMs independently to select top-$k$ models. However, this ignores model correlations and enforces a rigid computational budget. Consequently, routers often select redundant models that share failure modes, limiting the overall probability of success. To address this, we propose FlexRouter, a routing framework that explicitly models model complementarity. FlexRouter optimizes for answer coverage, maximizing the probability that at least one selected model yields a correct response. This objective aligns with practical inference pipelines where multiple candidate outputs are generated and a downstream verifier or user selects the final one. We formulate routing as a coverage-oriented subset selection problem and model the routing policy using Determinantal Point Processes (DPPs), which naturally capture both model competence and redundancy. To directly optimize coverage without requiring a ground-truth target subset, we introduce a training objective based on marginalizing over failure sets. During inference, we employ a greedy strategy based on marginal log-determinant gains, enabling the router to adaptively determine subset sizes without a predefined budget. Extensive experiments on the large-scale RouterEval
benchmark demonstrate that our proposed FlexRouter achieves higher coverage with lower redundancy across both in-domain and out-of-domain tasks than strong baselines while maintaining flexible inference cost.
Read the full abstract ▾
All authors Wang Wei, Harry Yang, Tiankai Yang, Samyadeep Basu, Hongjie Chen, Andy Zhao, Franck Dernoncourt, Ryan A. Rossi, Hoda Eldardiry
Machine Learning · Cross-listed from cs.AI
Conditional Generation of Creative Chess Puzzles with Diffusion Models
Aatu Selkee, Severi Rissanen, Xidong Feng et al. (+2 more)
While modern
language models demonstrate impressive generative capabilities, they often struggle with constrained, counter-intuitive creative tasks. To address this limitation, we explore chess puzzle generation as a rigorous testbed for computational creativity and
reasoning, a domain where altering a single piece can invalidate an entire solution. We propose a novel approach for conditional generation of creative chess puzzles using masked
diffusion models. Unlike previous methods, our non-directional
diffusion approach allows for conditioning on specific tactical themes and partial board positions. We introduce a novel auxiliary task of simultaneous best-move prediction, which improves solution uniqueness by 11.6% and theme-conditioning accuracy by 2.5%. To further optimize solution uniqueness and theme conditioning, we establish a reinforcement learning framework adapted from Denoising
Diffusion Policy Optimization (DDPO). This RL training increases the yield of unique and theme-matching positions by 89.1%. Finally, we release the first open-weights models (Appendix B) for chess puzzle generation, offering a new pathway for controllable, creative generation.
Read the full abstract ▾
All authors Aatu Selkee, Severi Rissanen, Xidong Feng, Tom Zahavy, Eric Malmi
Machine Learning
The Advantages of Fresh Sketching for Ridge Regression
Linkai Ma, Qilin Li and Petros Drineas
Over the past 25 years, sketching and sampling have become widely used tools for accelerating large-scale regression. In iterative randomized solvers, a basic design choice is whether to $\textit{reuse}$ the same sketch or draw $\textit{fresh}$ randomness at every step. For (under-constrained) iterative ridge regression with column sampling, whether fresh sketches offer provable advantages has remained open: $\textit{We show that they do.}$ Fresh sketching lets us analyze error only along the current residual solution, rather than uniformly over the entire Gram
matrix. This directional view yields sharper convergence guarantees for leverage score and ridge leverage score sampling and, more importantly, leads to residual-aware sampling rules. By minimizing the variance of the relevant sketched
matrix-vector product, we derive an oracle distribution and practical approximations to the oracle distribution, including a mixture sampling distribution with (somewhat weaker) convergence guarantees. Experiments on synthetic and real data, including ridge probes on Qwen2.5 representations, support our theory, showing substantially faster convergence.
Read the full abstract ▾
Machine Learning · Cross-listed from math.OC
Parameter-Free Zeroth-Order Optimization with Ellipsoidal Sampling
Alexander Sholokhov and Alexander Rogozin
Zeroth-order optimization methods are essential for solving black-box problems where gradient information is unavailable or expensive to compute. This paper presents POEM-ES, a novel parameter-free
stochastic zeroth-order
algorithm that extends the recent POEM method by integrating subspace preconditioning with ellipsoidal randomized sampling. In contrast to traditional zeroth-order approaches that rely on isotropic random directions, POEM-ES performs anisotropic sampling guided by a fixed structural symmetric positive semi-definite (SPSD) preconditioner $\hat{\Sigma}$ that encodes the underlying low-dimensional geometry. Under a standard structural spectral normalization where $\lambda_{\max}(\hat{\Sigma}) = 1$, we introduce the use of the empirical effective dimension $d^* = \operatorname{tr}(\hat{\Sigma})$, which reflects the intrinsic dimensionality of the problem and guides both the sampling and randomized smoothing parameter schedules. In practice, such a preconditioner can be effectively obtained via pilot sampling, historical trajectories, or domain-specific expert knowledge. We prove that POEM-ES achieves a dimension-reduced convergence rate under low-rank structural assumptions, requiring only $\tilde{\mathcal{O}}\left( \frac{\left( r^2 \kappa(\hat{\Sigma}) + d^* \right) L^2 D_{\mathcal{X}}^2}{\varepsilon^2} \right)$
stochastic zeroth-order oracle queries. The method remains fully parameter-free and demonstrates significant improvements over the original POEM in problems with low-rank structure where $d^* \ll d$. Numerical experiments on hinge-loss binary classification tasks using LibSVM
datasets confirm the practical superiority of the proposed approach.
Read the full abstract ▾
Machine Learning · Cross-listed from cs.HC
Fairness Theatre: Evaluating Post-Hoc Fairness Interventions in Vendor-Controlled Early Warning Systems
Kelly McConvey, Angelina Zhai, Rebecca Li et al. (+1 more)
Public institutions increasingly procure AI systems whose design they cannot inspect or change. In higher education, proprietary Early Warning Systems (EWS) leave colleges with few options beyond adjusting model outputs to address inequity. This raises the question of how
fairness work is coordinated among vendors, institutions, advisors, and students with unequal power to change these systems? Using student records from a public college in Ontario, Canada, we evaluate six post-hoc
fairness interventions on a research EWS under simulated procurement constraints. We compare
fairness, accuracy, and demographic disparities, introducing error-type profiling to trace how interventions redistribute false positives and false negatives. Interventions redistributed disparities without consistently reducing them. Two implementations favored already-advantaged groups because they used group size to define disadvantage; small, marginalized groups remained poorly served. These findings show how procurement constraints and implementation choices shape the possibilities for
fairness work. We call the resulting condition
fairness theatre; dashboard metrics converge while groups' error burdens persist or worsen.
Read the full abstract ▾
All authors Kelly McConvey, Angelina Zhai, Rebecca Li, Shion Guha