ICML 2026 Papers

The design of environments plays a critical role in shaping the development and evaluation of cooperative multi-agent reinforcement learning (MARL) algorithms. While existing benchmarks highlight critical challenges, they often lack the modularity required to design custom evaluation scenarios. We introduce the Totally Accelerated Battle Simulator in JAX (TABX), a high-throughput sandbox designed for reconfigurable multi-agent tasks. TABX provides granular control over environmental parameters, permitting a systematic investigation into emergent agent behaviors and algorithmic trade-offs across a diverse spectrum of task complexities. Leveraging JAX for hardware-accelerated execution on GPUs, TABX enables massive parallelization and significantly reduces computational overhead. By providing a fast, extensible, and easily customized framework, TABX facilitates the study of MARL agents in complex structured domains and serves as a scalable foundation for future research. Our code is available at: https://anonymous.4open.science/r/TABX-00CA.

The core challenge of Test-Time Reinforcement Learning (TTRL) lies in estimating rewards without access to ground-truth supervision. Existing TTRL methods predominantly rely on majority voting to generate pseudo-labels, under the assumption that the most frequent answer among sampled trajectories is correct. However, we observe that this assumption frequently breaks down in complex reasoning tasks, where correct solutions often constitute a logical minority. As a result, rare yet correct trajectories are systematically undervalued by majority-voting-based approaches. To address this limitation, we propose Self-Reflective Test-Time Reinforcement Learning (SR-TTRL), a novel framework that leverages self-reflective verification to produce high-fidelity pseudo-labels. Specifically, given multiple sampled trajectories for a problem, SR-TTRL first groups trajectories according to their final answers and selects one representative from each group to form a candidate pool. Each candidate trajectory is then summarized to preserve its core reasoning steps while reducing verbosity. Finally, the model performs self-reflection over the candidate pool, critically evaluating and selecting the most plausible trajectory as the pseudo-label. Empirically, SR-TTRL achieves substantially higher pseudo-label fidelity and sample efficiency than prior majority-voting-based TTRL methods. Extensive experiments across diverse benchmarks and model families demonstrate that SR-TTRL consistently outperforms majority-voting baselines and significantly improves generalization to novel problems. For example, SR-TTRL improves the Pass@1 accuracy of Qwen3-8B on AIME24 from 29.1 to 55.8 (a gain of +26.7), exceeding standard TTRL by an additional +9.1. The code will be released at: https://github.com/JIA-Lab-research/SR-TTRL.

Safe Reinforcement Learning (RL) algorithms are typically evaluated under fixed training conditions. We investigate whether training-time safety guarantees transfer to deployment under distribution shift, using diabetes management as a safety-critical testbed. We benchmark safe RL algorithms on a unified clinical simulator and reveal a safety generalization gap: policies satisfying constraints during training frequently violate safety requirements on unseen patients. We demonstrate that test-time shielding, which filters unsafe actions using learned dynamics models, effectively restores safety across algorithms and patient populations. Across eight safe RL algorithms, three diabetes types, and three age groups, shielding achieves Time-in-Range gains of 13--14\% for strong baselines such as PPO-Lag and CPO while reducing clinical risk index and glucose variability. Our simulator and benchmark provide a platform for studying safety under distribution shift in safety-critical control domains. Code is available at https://github.com/safe-autonomy-lab/GlucoSim and https://github.com/safe-autonomy-lab/GlucoAlg.

Frontier language model agents can exhibit misaligned behaviors, including deception, exploiting reward hacks, and pursuing hidden objectives. To control such agents, we can use LLMs themselves to *monitor* for misbehavior. In this paper, we study how *information access* affects LLM monitor performance and find that contemporary systems often perform better with less information, a phenomenon we call the *less-is-more effect* for automated oversight. We demonstrate this phenomenon, and analyze the conditions under which it occurs, in three evaluation environments where agents must conduct sabotage while evading monitors. Leveraging this effect, we introduce *extract-and-evaluate* (EaE) monitoring—a new hierarchical approach where one monitor isolates relevant excerpts from the monitored agent's trajectory and another monitor scores these excerpts. In BigCodeBench-Sabotage with a GPT-4.1-mini monitor, EaE improves sabotage detection rates by 16.8 percentage points over the next-best approach. In other settings, EaE either outperforms or is competitive with baselines. Additionally, we find that agents unaware of being monitored can be caught much more easily, that monitors scoring longer reasoning traces gain more from information filtering, and that monitor performance scales linearly with cost.

Finding flat minima in the loss landscape is a key strategy for Domain Generalization (DG). However, its effectiveness is often limited by two crucial challenges. 1) Domain Shift: Existing methods like Sharpness-Aware Minimization (SAM) apply a uniform optimization strategy across all domains, overlooking the differences of the learning difficulties among multiple domains and thus performing poorly on challenging domains. 2) Anisotropic Sharpness: By perturbing parameters along a single gradient direction, SAM and its variants ignore multi-directional flatness, making the model converge to minima that remain sharp in other directions. The combined challenges make it more difficult for the model to find truly robust solutions in multi-domain scenarios. To overcome these limitations, we propose the Dual Adaptive Sharpness-Aware Minimization (DA-SAM), which comprises two key modules: Dynamic Adaptive Scaling (DAS) module and Adaptive Multi-Directional Flattening (AMDF) module. First, to tackle the domain shift problem, the DAS module computes the real-time loss on each domain to adaptively generate domain-specific scaling factors that guide the generation of perturbation directions. Second, the AMDF module calculates local flatness by generating multiple directions to simulate perturbations in the parameter space. Based on the learned local flatness metric, it dynamically adjusts the perturbation step size to guide the model parameters to be away from anisotropic sharp regions. Crucially, DAS provides domain-level guidance that makes AMDF’s multi-directional geometric exploration more targeted and effective. Extensive experiments on five DG benchmarks demonstrate the effectiveness of our DA-SAM algorithm.

Robust machine learning for graph-structured data has made significant progress against test-time attacks, yet certified robustness to poisoning – where adversaries manipulate the training data – remains largely underexplored. For image data, state-of-the-art poisoning certificates rely on partitioning-and-aggregation schemes. However, we show that these methods fail when applied in the graph domain due to the inherent label and structure sparsity found in common graph datasets, making effective graph-partitioning difficult. To address this challenge, we propose a novel semi-supervised learning framework called deep Self-Training Graph Partition Aggregation (ST-GPA), which enriches each graph partition with informative pseudo-labels and synthetic edges, enabling effective certification against node-label and graph-structure poisoning under sparse conditions. Our method is architecture-agnostic, scales to large numbers of partitions, and consistently and significantly improves robustness guarantees against both label and structure poisoning across multiple benchmarks, while maintaining strong clean accuracy. Overall, our results establish a promising direction for certifiably robust learning on graph-structured data against poisoning under sparse conditions.

Reinforcement Learning from Human Feedback (RLHF) can reveal implicit objectives such as safety considerations that go beyond task completion. In this work, we focus on the common safety criteria embedded in crowd preference datasets, where different users may express distinct preferences or objectives, yet follow similar safety principles. Our aim is to discover shared safety criteria from crowd preferences and then transfer them to downstream RL tasks to regularize agent behavior and enforce safety. We first show that direct reward combination—optimizing a preference-learned reward model together with downstream task rewards—has inherent limitations. Motivated by this, we propose Safe Crowd Preference-based RL, a hierarchical framework that extracts safety-aligned skills from crowd preferences and composes them via a high-level policy to safely solve downstream tasks. Experiments across safe RL environments and a preliminary LLM-style task with diverse user goals and shared safety constraints demonstrate that our approach substantially lowers safety costs without access to explicit safety rewards, while achieving task performance comparable to oracle methods trained with ground-truth safety signals.

Driven by the imperative to leverage privacy-sensitive data scattered across decentralized devices, federated fine-tuning has emerged as a vital paradigm for adapting large language models without compromising data privacy. Yet, its practical efficacy is bottlenecked by severe client resource heterogeneity. Existing truncation-based methods typically couple the transmitted rank with the trainable rank, which (i) under-utilizes bandwidth on communication-rich but compute-limited clients and (ii) exacerbates truncation-induced gradient drift. To address this, we propose FedHera, a resource-decoupled framework that explicitly differentiates information reception from gradient optimization. FedHera employs a spectrum-preserving allocation strategy to maximize the transfer of global knowledge (via high-rank singular values) within bandwidth limits, irrespective of training constraints. Furthermore, we introduce a prefix-gating mechanism that utilizes the downloaded high-capacity basis as a frozen reference to guide local updates, thereby minimizing the optimization gap caused by aggressive truncation. Extensive experiments under different heterogeneous settings show that FedHera improves stability and accuracy over state-of-the-art baselines.

Millimeter-wave (mmWave) sensing enables privacy-preserving, always-on edge perception, but its measurements are often sparse, temporally irregular, and corrupted by high-frequency noise. Existing mmWave pipelines predominantly rely on artificial neural networks (ANNs), which achieve robustness through extensive preprocessing or deep architectures, thereby limiting their efficiency on edge devices. In this work, we study spiking neural networks (SNNs) for mmWave sensing from a mechanism–data alignment perspective. By leveraging the low-pass filtering behavior of leaky integrate-and-fire (LIF) dynamics, we analyze how their implicit temporal filtering interacts with the frequency structure of mmWave signals. Our analysis shows that when discriminative information resides in low-to-mid frequencies, LIF dynamics can inherently suppress high-frequency noise, clarifying when and why SNNs outperform ANNs. Based on this insight, we derive a principled criterion for configuring the membrane decay factor by matching the effective bandwidth of LIF dynamics to the data’s discriminative spectral content. Experimental results across four widely used mmWave datasets validate the proposed frequency-matching hypothesis, yielding an average test-accuracy improvement of 6.22% and a 3.64× reduction in theoretical energy consumption relative to ANN baselines, under a unified evaluation protocol.

Flow Matching models have demonstrated strong performance across a wide range of generative tasks. However, their reliance on ODE-based iterative sampling incurs substantial computational overhead, which limits their applicability in real-time scenes. While distillation is a promising solution, existing approaches largely borrow from diffusion-based score matching, often failing to exploit the intrinsic geometric structure of flows and suffering from training instability, high variance, and degraded generation quality. In this paper, we propose Mean Flow Distillation (MFD), a novel distillation framework tailored for flow matching models. We theoretically demonstrate that MFD acts as a temporal low-pass filter, effectively suppressing the high-frequency optimization noise inherent in variational score distillation (VSD) while ensuring global trajectory consistency. We further prove the Mean Flow Matching Theorem, establishing that matching expected average velocities is sufficient for strict distribution alignment. Empirically, on challenging high-dimensional manifolds including 4D occupancy forecasting and text-to-image generation, MFD achieves state-of-the-art performance, enabling high-fidelity single-step generation.

Step-level caching accelerates diffusion models by exploiting temporal redundancy across denoising steps. Existing methods make per-step cache decisions using threshold-based heuristics, without directly optimizing for final output quality. As a result, their inference latency varies across inputs and is difficult to control at deployment. In this work, we propose BudCache, which inverts this formulation: rather than letting per-step error thresholds dictate the runtime cost, we fix the compute budget in advance and search for the cache policy that best preserves the final output. To tackle the combinatorial complexity of step selection, we combine Simulated Annealing with deterministic Hill Climbing. This offline search identifies high-quality cache policies within minutes and introduces no online search or thresholding overhead during inference. When the compute budget is very tight, we further introduce cache-aware schedule alignment, which adapts the time discretization to the selected cache policy to reduce cache-induced trajectory mismatch. Experiments on FLUX.1-dev and Wan2.1 show that BudCache achieves better generation quality than heuristic caching baselines under the same inference budgets.

Selective State Space Models (SSMs) such as Mamba have emerged as efficient alternatives to Transformers, achieving linear complexity through input-dependent parameterization. However, this selectivity transforms the system from linear time-invariant (LTI) to linear parameter-varying (LPV), where individually stable matrices can produce unbounded trajectories under switching. Existing work focuses on empirical performance, leaving global stability, robustness bounds, and practical certification unresolved. This paper develops a control-theoretic framework providing the first comprehensive stability and robustness analysis for selective SSMs. We prove BIBO stability by viewing selective scans as continuous-time LTI sampling and establish two-term robustness bounds with linear growth in sequence length. For general LPV systems, we provide common quadratic Lyapunov function conditions and develop algorithms to extract certificate constants directly from trained weights. These results bridge control theory and SSM architectures, enabling formal guarantees for safety-critical deployment.

Deep neural networks often generalize well despite heavy over-parameterization, challenging classical parameter-based analyses. We study generalization from a representation-centric perspective and analyze how the geometry of learned embeddings is associated with generalization performance for a fixed trained model. We derive a post-hoc generalization bound that relates the gap between population risk and held-out empirical risk to two factors: (i) the intrinsic dimension of the embedding, which determines the convergence rate of the empirical embedding distribution to its population counterpart in Wasserstein distance, and (ii) the sensitivity of the downstream mapping from embeddings to predictions, characterized by Lipschitz constants. Together, these provide a post-hoc explanation of generalization for trained models. At the final embedding layer, architectural sensitivity disappears and the bound is dominated by embedding dimension, explaining its strong empirical correlation with generalization performance. Experiments across architectures and datasets validate the theory and demonstrate the utility of embedding-based diagnostics.

Code repair is an important capability for language models (LMs): given a buggy program and unit tests, an LM must produce a fixed program that passes the tests. We aim to scale supervision for code repair by having an LM generate bug--fix tasks with unconstrained edits, using unit tests as the only verifier. We propose generator-fixer self-play, in which a single model is trained with reinforcement learning to alternate between generating bugs and fixing them. As the fixer improves, the generator adapts to produce increasingly difficult bugs, yielding an automatic curriculum. However, because unit tests certify correctness but not realism, we find that the generator can drift from bugs encountered in practice, improving repair on self-generated bugs while degrading on real-world bugs. We propose Anchored Self Play (ASP), which anchors self-play with a small reference set by (i) adding a code-embedding similarity reward to guide generation and (ii) mixing reference bugs into fixer training to prevent drift. To reflect LM-assisted programming, where bugs come from humans, LMs, and human edits of LM code, we introduce BugSourceBench, a code repair benchmark spanning human-authored bugs, human-edited buggy LM code, and errors in LM-generated code. Across bug sources, ASP achieves the best fix rates, improving average fix rate by $+25$% (relative) / $+7.2$ pp (absolute) over standard self-play, with gains on both LM-error bugs ($+100$% relative / $+11$ pp absolute) and human-authored bugs ($+7.1$% relative / $+3.4$ pp absolute).

Though biological foundation models (Bio-FMs) have delivered strong performance across biomedical tasks, their robustness to small-but-real perturbations is underexplored. In this work, we ask: Are Bio-FMs robust for real-world use? What perturbations compromise their reliability? Our pilot study suggests that due to subtle biological data curation issues and common machine-learning (ML) processing choices, Bio-FMs suffer from two complementary perturbation sources: biologically plausible perturbations (capturing experimental corruptions and curation artifacts) and ML-induced transformations (capturing preprocessing, data augmentation, and embedding choices). Guided by this taxonomy, we design perturbation suites that mimic corruptions frequently encountered in biological experiments, and we systematically probe how transformations in the ML pipeline reshape model behavior. By conducting 2,128 experiments over 11 state-of-the-art Bio-FMs on 7 bio-tasks, we show that most Bio-FMs are vulnerable to both biological perturbations and ML transformations, revealing underappreciated robustness gaps that can directly translate into deployment risk. Interestingly, we find that subtle biological perturbations, which are often imperceptible to current measurement tools, can induce severe discrepancies in Bio-FM outputs and lead to critical failures, yet cryo-EM models (e.g., CryoDRGN) exhibit a surprising level of robustness even under worst-case perturbations. Our study for the first time surfaces critical failure modes and provides a principled perspective for evaluating the robustness of Bio-FMs.

Recent leaderboard-based evaluations of large language models aggregate user feedback by fitting a Bradley--Terry model to pairwise comparisons, producing a single global ranking based on a latent quality score. While appealing for its simplicity, this approach is incompatible with heterogeneous preferences: when LLMs are used across diverse tasks and use cases, users who favor fundamentally different model behaviors can be systematically misrepresented when collapsed into a single quality score. To address this issue, we study *pluralistic leaderboards* that aim to remain *stable* with respect to heterogeneous user populations. Drawing on ideas from social choice theory, we adapt the notion of *local stability*, which requires that no model outside the top-$k$ positions is collectively preferred to the top-$k$ set by more than $O(1/k)$ fraction of users. Building on techniques from the social choice literature, we design an alternative leaderboard mechanism that satisfies local stability while eliciting only $\widetilde{O}(k)$ pairwise comparisons per user, where $k$ is the size of the prefix for which stability is guaranteed. Using data from LMArena, we show that standard Bradley--Terry aggregation can violate local stability in practice, whereas our method provides substantially stronger stability guarantees.

The rapid expansion of remote sensing technology has generated massive amounts of unlabeled multi-view data distributed across different institutions. Analyzing this data presents significant challenges, as centralized processing incurs prohibitive communication costs and raises data privacy concerns. To address these issues, this paper proposes a novel deep federated multi-view clustering (MVC) framework tailored for remote sensing data. Unlike existing methods that transmit sensitive data features, our approach shares only privatized prototypes masked with adaptive noise, ensuring both communication efficiency and privacy preservation. First, we employ superpixel segmentation to reduce the spatial dimensionality of remote sensing data, lowering computational burdens. Furthermore, to resolve the inconsistency of cluster assignments across different clients, we design a co-occurrence structural alignment module that synchronizes local models. Finally, we incorporate a wasserstein prototype contrastive learning mechanism, which models clusters as distributions rather than points, to enhance global consistency and robustness against data heterogeneity. Extensive experiments on four public datasets demonstrate that our framework achieves superior clustering performance and efficiency compared to state-of-the-art methods.

Low-Rank Adaptation (LoRA) is the most widely adopted method for fine-tuning large language models. Notably, LoRA is inherently overparameterized: multiple pairs of low-rank factors can yield the same adapted weight matrix. We show—both theoretically and empirically—that these pairs exhibit significantly different condition numbers. As a result, converging to different loss minimizers directly impacts the convergence rate of LoRA. Building on this observation, we introduce Balanced Low-Rank Adaptation (BaLoRA), a variant of LoRA that projects iterates onto a balanced manifold. This manifold improves the conditioning of the loss landscape while preserving the adapted matrix. The projection step is computationally lightweight and integrates seamlessly into existing fine-tuning pipelines. Empirically, BaLoRA converges faster than standard LoRA and achieves superior performance across a range of fine-tuning tasks.

Latent or continuous chain-of-thought methods replace explicit textual rationales with a number of internal latent steps, but these intermediate computations are difficult to evaluate beyond correlation-based probes. In this paper, we view latent chain-of-thought as a manipulable causal process in representation space by modeling latent steps as variables in a structural causal model (SCM) and analyzing their effects through step-wise $\mathrm{do}$-interventions. We study two representative paradigms (i.e., Coconut and CODI) on both mathematical and general reasoning tasks to investigate three key questions: (1) which steps are causally necessary for correctness and when answers become decodable early; (2) how influence propagates across steps and how this structure compares to explicit CoT; and (3) whether intermediate trajectories retain competing answer modes and how output-level commitment differs from representational commitment across steps. We find that latent-step budgets behave less like homogeneous extra depth and more like staged functionality with non-local routing, and we identify a persistent gap between early output bias and late representational commitment. These results motivate mode-conditional and stability-aware analyses, together with corresponding training/decoding objectives, as more reliable tools for interpreting and improving latent reasoning systems. Code is available at https://github.com/J1mL1/causal-latent-cot.