ICML 2026 Papers

Learning and planning in imagination using world models provides an effective paradigm for training agents for decision-making. However, existing approaches often rely on high-dimensional latent spaces or generic visual embeddings that retain many factors irrelevant to control, limiting efficiency and generalization across tasks. To this end, we study how agents can learn world models with representations that are task-specific, minimal, and sufficient for decision making. We achieve this via a closed-loop synergy between the agent and the world model, in which structured world-model learning distills task-sufficient representations from informative interaction data. On the agent side, agents perform active probing of the environment to collect informative trajectories that expose task-relevant latent factors, guided by an adaptive curriculum. On the world-model side, we learn structured representations over observations to distill compact, task-sufficient latent states from the collected interaction data. This synergy enables the recovery of task-sufficient latent representations that capture all control-relevant factors empirically. Leveraging these representations, the resulting policies achieve improved sample efficiency generalization, including generalization across skills, object–skill compositions, and previously unseen tasks on standard continuous control and robotic manipulation benchmarks.

Text-to-image diffusion models have achieved remarkable success in generating high-quality images, yet existing safety mechanisms exhibit critical cross-seed instability where defense performance varies significantly under different random seed conditions. This instability stems from the fact that a single malicious prompt generates diverse harmful variants across different noise initializations, forming complex distributional clusters that current methods cannot adequately address. We investigate extending Noise Contrastive Alignment (NCA) to diffusion models due to its native capability of handling multiple negative samples through probabilistic weighting, but our theoretical analysis reveals two fundamental flaws in direct extension: gradient reversal caused by positive regularization terms that paradoxically penalize safe content generation, and uniform suppression of harmful samples that ignores severity variations. To tackle these issues, we propose Noise Contrastive Diffusion (NCD), which incorporates targeted algorithmic modifications including elimination of problematic regularization and introduction of pairwise regularization mechanisms that establish individualized preference relationships between safe and harmful variants. Extensive experiments further demonstrate that NCD achieves superior cross-seed stability, reducing attack success rates (ASRs) from 11.1% to 6.2% compared to SOTA methods at the seed level while maintaining exceptional generation quality, exhibiting robust resistance against sophisticated jailbreak prompts and strong generalizability across different T2I architectures. WARNING: This paper may contain examples of harmful texts and images.

Reliable uncertainty quantification at unobserved spatial locations is a key challenge in spatial statistics, particularly for complex and heterogeneous datasets. While traditional methods such as Kriging rely on strong distributional assumptions, conformal prediction (CP) offers a distribution-free alternative. However, although non-i.i.d. CP theory is well established for time-series data, a significant gap remains for spatial data, where the lack of a natural ordering and discrete index complicates theoretical guarantees. Existing CP theory for spatial data often relies on exchangeability. We propose Localized Spatial Conformal Prediction (LSCP), a model-agnostic framework that bridges this gap by coupling local quantile regression with conformal calibration. LSCP conditions on spatial neighborhoods to capture local heterogeneity. We show that LSCP retains finite-sample marginal coverage under spatial exchangeability and attains asymptotic conditional coverage under stationarity and spatial mixing. Across synthetic and real datasets, LSCP consistently achieves near-nominal coverage with tighter and more stable prediction intervals than existing methods that fail to capture these spatial dependencies.

Online convex optimization in non-stationary environments has garnered considerable attention in the literature. Recently, Pasteris et al. (2024) investigate online convex optimization with the optimal switching regret, defined as the sum of the static regret over each segment, where the segmentation is an arbitrary partition of the entire time horizon. For general convex functions, their work has established an optimal switching regret bound. However, it remains open whether similar bounds are attainable for other types of convex functions, such as exponentially concave or strongly convex functions. In this paper, we affirmatively answer this question by proposing a novel meta-algorithm, termed IRESET, which is used to aggregate the decisions from a group of experts. The essence of our method lies in running multiple experts over a set of intervals, and then employing a meta-algorithm equipped with second-order bounds to sequentially combine their decisions. We leverage the segment tree structure to analyze the switching regret over the entire time horizon, and offer new insights into utilizing recursive equations over the segment tree. By choosing appropriate expert-algorithms for IRESET, our methods achieve logarithmic switching regret bounds for exponentially concave or strongly convex functions, respectively.

While AI models often refuse explicitly unlawful requests, in real-world scenarios illegality often depends on context. We evaluate frontier models on contextual illegality across four corporate law scenarios in which routine actions—editing documents, trading stock, requesting payment, approving communications—become unlawful due to circumstances such as pending investigations or bankruptcy filings. We study both chat and agentic settings and compare results to a human baseline. The best-performing models consistently followed lawful requests and refused unlawful requests, though performance varied substantially between different scenarios and models. We also identify distinct failure modes, such as excessive refusal of lawful requests, and find higher performance in reasoning models and agentic environments. By studying contextual illegality in these controlled environments, we develop a methodology that can be extended to evaluate the legal compliance of AI models in additional scenarios and domains.

Hypergraph Neural Networks (HNNs) model high-order interactions effectively but rely on costly node annotations, motivating Hypergraph Active Learning (HAL). However, many HAL pipelines adapt graph-based querying through clique expansion, which introduces *structural bias* and can cause *ranking collapse*, making utilities overly determined by hyperedge cardinalities rather than informative high-order context. We propose **HIAL** (**H**ypergraph **I**nfluence-based **A**ctive **L**earning), a training-free framework that formulates hypergraph active learning as influence maximization over a high-order context-based weighted pairwise projection of the hypergraph. HIAL employs a High-Order Interaction-aware propagation mechanism that modulates pairwise influence weights using hyperedge cardinality and feature consistency, yielding a scalable linear diffusion process tailored to homophilic hypergraphs. We further combine feature-space coverage and structural reachability into a monotone submodular selection objective, enabling an efficient lazy greedy solver. Experiments on eight benchmarks demonstrate that HIAL consistently outperforms strong baselines across diverse homophilic hypergraph domains.

High-dimensional genomics studies are frequently confounded by unmeasured biological processes that obscure disease-specific signals. While existing workflows can estimate these latent confounders, they fail to quantify how robust a discovery is to varying levels of hypothetical confounding. We introduce sensGAN, a deep-learning adversarial framework that systematically explores the confounding spectrum by learning "worst-case" latent variables that nullify the most gene associations under novel predictive-gain constraints. By identifying the minimum confounding strength required to explain away an observed effect, our method shifts the paradigm toward a formal, quantitative sensitivity analysis. In diverse simulations, sensGAN accurately recovers latent structures and outperforms existing methods in identifying confounder-sensitive genes. Applied to human Alzheimer's disease microglia, our framework prioritizes robust disease pathways while successfully isolating signals driven by unmeasured co-occurring neurodegenerative pathologies.

The fields of AI-based disease fingerprinting, drug discovery and repurposing are currently among the emerging frontiers of machine learning applied to medicine. One major challenge is to obtain robust $\textit{in-silico}$ modeling of disease progression while accounting for the vastly different time scales of biochemical interactions, from gene expression to protein abundance and metabolic flux. Discrete sequence models inadequately represent such multi-scale interactions, and standard Neural Ordinary Differential Equations (NODEs) often fail to train stably under stiffness (different time scales). To address this, a Tri-Scale Stiff NODE, defined by hierarchically coupled latent differential equations that model the causal relationships from genes to proteins and metabolites, is introduced and optimized in this paper in terms of reconstruction error and information-theoretic mutual information. This enables continuous-time modeling of cellular responses to identify not only the disease dynamics, but also drug perturbations that act within narrow time windows, often invisible to discrete-time approaches. Lyapunov analysis provides a theoretical guarantee that the modeled trajectories remain stable and well-behaved even under extreme stiffness. The methodology is validated using the STATegra B-cell and Traxler macrophage datasets, with the former utilized for a proof-of-concept drug repurposing pipeline.

Associative memory has long underpinned the design of sequential models. Beyond recall, humans reason by *projecting future states and selecting goal-directed actions*, a capability that modern language models increasingly require but do not natively encode. While prior work uses reinforcement learning or test-time training, planning remains external to the model architecture. We formulate reasoning as *optimal control* and introduce the *Test-Time Control (TTC)* layer, which performs finite-horizon *LQR* planning over latent states at inference time, represents a *value function* within neural architectures, and leverages it as the nested objective to enable *planning before prediction*. To ensure scalability, we derive a hardware-efficient LQR solver based on a symplectic formulation and implement it as a fused CUDA kernel, enabling parallel execution with minimal overhead. Integrated as an adapter into pretrained LLMs, TTC layers improve mathematical reasoning performance by up to +27.8% on MATH-500 and 2-3$\times$ Pass@8 improvements on AMC and AIME, demonstrating that embedding optimal control as an architectural component provides an effective and scalable mechanism for reasoning beyond test-time training.

Distributional reinforcement learning (DRL) models the full return distribution, but typically relies on finite-dimensional categorical or quantile approximations, often involving projection or quantile-regression approximations to the Bellman target, together with independently sampled bootstrap targets that obscure transport structure and add variance. We present Path-Coupled Bellman Flows (PCBF), a continuous-time DRL method that encodes Bellman endpoint consistency and pathwise Bellman-coupled geometry within generative flow trajectories. PCBF represents return distributions via flow matching and couples the paths of consecutive states through shared base noise, yielding a geometric Bellman flow relation between velocity fields. This structure enables a $\lambda$-parameterized control-variate target that reduces training variance while preserving the source and Bellman endpoint geometry. Experiments on analytically tractable MRPs, OGBench, and D4RL show improved distributional fidelity, training stability, and competitive offline RL performance.

Symmetries are known to improve the empirical performance of machine learning models, yet theoretical guarantees explaining these gains remain limited. Prior work has focused mainly on compact group symmetries and often assumes that the data distribution itself is invariant, an assumption rarely satisfied in real-world applications. In this work, we extend generalization guarantees to the broader setting of non-compact symmetries, such as translations and to non-invariant data distributions. Building on the PAC-Bayes framework, we adapt and tighten existing bounds, demonstrating the approach on McAllester's PAC-Bayes bound while showing that it applies to a wide range of PAC-Bayes bounds. We validate our theory with experiments on several datasets with non-uniform and non-compact transformations, where the derived guarantees not only hold but also improve upon prior results. These findings provide theoretical evidence that, for symmetric data, symmetric models are preferable beyond the narrow setting of compact groups and invariant distributions, opening the way to a more general understanding of symmetries in machine learning.

Despite remarkable progress in Multimodal Large Language Models (MLLMs), these models still struggle with fine-grained understanding tasks. In this work, we propose **Procedurally Generated Tasks (PGT)** a simple data-driven framework that serves a dual purpose: inducing fine-grained visual understanding and acting as a low-cost diagnostic tool to identify the source of perception failures. By overlaying unambiguous geometric primitives on images, PGT generate additional dense supervision that disentangles visual grounding capability from semantic priors. Extensive experiments on relational, quantitative, and 3D/depth understanding benchmarks show that PGT yields remarkable gains across diverse architectures. Instruction tuning MLLMs on LLaVA-v1.5-Instruct augmented with PGT data results in improvements of up to +20\% on the What’sUp benchmark and +13.3\% on CV-Bench-2D, while maintaining general perception capabilities. Moreover, finetuning state-of-the-art MLLMs on PGT data leads to boosts of up to +5.5\% on What’sUp and +8.3\% on CV-Bench-2D. These findings demonstrate that PGT effectively address the bottleneck of fine-grained perception, revealing that many spatial reasoning deficits stem from inadequate supervision signals rather than inherent architectural or resolution limitations.

Contemporary AI faces the challenge of balancing generality with user-specific personalization. In federated learning (FL), this challenge is amplified by highly heterogeneous client data with complex non-IID patterns beyond standard modeling assumptions. Many existing FL methods are designed for relatively restricted heterogeneity settings (e.g., a fixed number of clusters or a fixed form of personalization), limiting their robustness under complex structures. In this work, we study FL from a \emph{multi-level non-IID} perspective, where client similarity is approximated by multiple granularities of shared knowledge: global, subgroup, and client-specific components. This view captures coarse-to-fine relationships while requiring less prior knowledge of task boundaries. Building on this insight, we propose \emph{Federated Multi-level Additive Modeling} (FeMAM), which learns multiple levels of shareable models and constructs personalized predictors via additive composition across levels. To move beyond a fixed structure, FeMAM allows models to grow and be pruned dynamically during training, adapting to diverse federated scenarios. Despite employing multiple models, FeMAM remains cost-friendly by activating only a small subset (one level) of models for training at a time. Extensive experiments show that FeMAM effectively approximates complex non-IID structures and consistently outperforms representative clustered and personalized FL baselines.

Most multimodal learning assumes redundant views (such as image–caption pairs), yet many applications require combining complementary modalities that provide distinct evidence (such as an X-ray and medical history). We term this setting *Complementary Multimodal Classification* (CMC). In CMC, existing explainable-by-design methods often force an accuracy–interpretability trade-off because single shared similarity metrics fail under asymmetric, class-conditional evidence. To address this, we propose Coherent Alignment of Multimodal Prototypes (CAMP). CAMP enforces coherent multimodal reasoning by aligning class-wise evidence via optimal transport and imposing geometric constraints to counter modality dominance and representation collapse. We provide theoretical guarantees showing that these mechanisms eliminate such degeneracies without restricting expressivity. Empirically, across 16 public CMC datasets, CAMP matches or exceeds large ($>$100M parameter) AutoML baselines with fewer than 1M trainable parameters, and when fine-tuned end-to-end it achieves state-of-the-art performance. To the best of our knowledge, this work is the first modality-agnostic prototype-learning framework designed for complementary multimodal tasks.

Transformer-based large language models face severe scalability challenges in long-context generation due to the computational and memory costs of full-context attention. Under practical computation and memory constraints, many inference-efficient long-context methods improve efficiency by adopting bounded-context or segment-level execution only during inference, while continuing to train models under full-context attention, resulting in a mismatch between training and inference execution and state-transition semantics. Based on this insight, we propose a training-consistent segment-level generation framework, in which training and inference follow the same segment-level forward execution semantics. During training, consistency with inference is enforced by restricting gradient propagation to KV states carried over from the immediately preceding segment, while permitting head-specific access to past KV states during the forward pass without involving them in gradient propagation. Across long-context benchmarks, our approach achieves performance comparable to full-context attention, while achieving competitive latency--memory trade-offs against strong inference-efficient baselines, and substantially improving scalability at very long context lengths (e.g., approximately $6\times$ lower peak prefill memory at 128K compared to full-context attention with FlashAttention).

Computation overhead has emerged as a critical bottleneck in Federated Learning (FL). Coreset selection tackles this challenge by constructing an informative subset to represent the full dataset. However, existing approaches optimize coreset construction solely at the data level and enforce a uniform retention ratio across all clients, ignoring client heterogeneity and introducing detrimental fragmented clients. In this paper, we first introduce a *gradient projection Shapley value* (GPSV) to evaluate client contributions. GPSV captures both the directional and magnitude information of client updates and enables exact Shapley value calculation with $\mathcal{O}(1)$ per-coalition evaluation. Building on GPSV, we then propose ShapCCS, the first client-level coreset selection strategy for FL. ShapCCS prioritizes clients with high GPSV scores while excluding fragmented clients with negligible or even negative GPSV. As a client-level coreset selection strategy, ShapCCS can be integrated with a data-level selection approach, and additionally reduces communication costs, an advantage unattainable by data-level methods alone. Extensive experiments demonstrate the superiority of ShapCCS on model performance and robustness to noise. The code is available at https://anonymous.4open.science/r/ShapCCS-5CBB.

Multimodal generative models require a unified approach to handle both discrete data (e.g., text and code) and continuous data (e.g., image, audio, video). In this work, we propose Latent Language Modeling (LatentLM), which seamlessly integrates continuous and discrete data using causal Transformers. Specifically, we employ a variational autoencoder (VAE) to represent continuous data as latent vectors and introduce next-token diffusion for autoregressive generation of these vectors. Additionally, we develop $\sigma$-VAE to address the challenges of variance collapse, which is crucial for autoregressive modeling. Extensive experiments demonstrate the effectiveness of LatentLM across various modalities. In image generation, LatentLM sis competitive with or outperforms DiT-style baselines under matched unified settings. When integrated into multimodal large language models, LatentLM provides a general-purpose interface that unifies multimodal generation and understanding. Experimental results show that LatentLM achieves favorable performance compared to Transfusion and vector quantized models in the setting of scaling up training tokens. In text-to-speech synthesis, LatentLM outperforms the state-of-the-art VALL-E 2 model in speaker similarity and robustness, while requiring 10 fewer decoding steps. The results establish LatentLM as a highly effective and scalable approach to advance large multimodal models.

Large language models (LLMs) have shown promise as interactive agents that solve tasks through extended sequences of environment interactions. While prior work has primarily focused on system-level optimizations or algorithmic improvements, the role of task horizon length in shaping training dynamics remains poorly understood. In this work, we present a systematic empirical study that examines horizon length through controlled task constructions. Specifically, we construct controlled tasks in which agents face identical decision rules and reasoning structures, but differ only in the length of action sequences required for successful completion. Our results reveal that increasing horizon length alone constitutes a training bottleneck, inducing severe training instability driven by exploration difficulties and credit assignment challenges. We demonstrate that horizon reduction is a key principle to address this limitation, stabilizing training and achieving better performance in long-horizon tasks. Moreover, we find that horizon reduction is related to stronger generalization across horizon lengths: models trained under reduced horizons generalize more effectively to longer-horizon variants at inference time, a phenomenon we refer to as horizon generalization.

Distributed stochastic optimization intertwines (i) stochastic gradient noise, (ii) communication compression, and (iii) adaptive/normalized updates. While each factor has been studied in isolation, their joint effect under realistic assumptions remains poorly understood. In this work, we develop a unified theoretical framework for Distributed Compressed SGD (DCSGD) and its sign variant Distributed SignSGD (DSignSGD) under the recently introduced $(L_0, L_1)$-smoothness condition. From a conceptual perspective, we show that the first- and second-order modified equations from the literature do not accurately model the discrete-time step-size/stability restrictions, especially under $(L_0,L_1)$-smoothness. From a technical perspective, we propose new first-order SDEs by carefully incorporating curvature-dependent terms into their drift: This helps capture the fine-grained relationship between learning rate restrictions, gradient noise, compression, and the geometry of the loss landscape. Importantly, we do so under general gradient noise assumptions, including heavy-tailed and affine-variance regimes, which extend beyond the classical bounded-variance setting. Our results suggest that normalizing the updates of DCSGD emerges as a natural condition for stability, with the degree of normalization precisely determined by the gradient noise structure, the landscape’s regularity, and the compression rate. In contrast, DSignSGD converges even under heavy-tailed noise with standard learning rate schedules. Together, these findings offer both new theoretical insights and perspectives, and practical guidance.