ICML 2026 Papers

Masked Diffusion Language Models have recently emerged as a powerful generative paradigm, yet their generalization properties remain understudied compared to their auto-regressive counterparts. In this work, we investigate these properties within the setting of the $k$-parity problem (computing the XOR sum of $k$ relevant bits), where neural networks typically exhibit grokking—a prolonged plateau of chance-level performance followed by sudden generalization. We theoretically decompose the Masked Diffusion (MD) objective into a Signal regime which drives feature learning, and a Noise regime which serves as an implicit regularizer. By training nanoGPT using MD objective on the $k$-parity problem, we demonstrate that MD objective fundamentally alters the learning landscape, enabling rapid and simultaneous generalization without experiencing grokking. Furthermore, we leverage our theoretical insights to optimize the distribution of the mask probability in the MD objective. Our method significantly improves perplexity for 50M-parameter models and achieves superior results across both pre-training from scratch and supervised fine-tuning. Specifically, we observe performance gains peaking at $8.8$% and $5.8$%, respectively, on 8B-parameter models, confirming the scalability and effectiveness of our framework in large-scale masked diffusion language model regimes.

While algorithmic stability is a central tool for understanding generalization of learning algorithms, existing high-probability guarantees typically rely on uniform boundedness or sub-Gaussian/sub-Weibull tail assumptions, which can be overly restrictive for modern settings with heavy-tailed or unbounded losses. We develop a stability-based framework that requires only a finite $L_p$ moment condition. Our first contribution is sharp concentration inequalities for functions of independent random variables under $L_p$ constraints, extending McDiarmid's bounded-differences techniques beyond the classical regime. Leveraging these results, we derive sharp high-probability generalization bounds across a range of learning paradigms, including empirical risk minimization, transductive regression, and meta-learning. These guarantees show that $L_p$ stability suffices for robust generalization even when boundedness fails, substantially weakening the standard assumptions in the stability literature.

We aim to develop physics foundation models for science and engineering that provide real-time solutions to Partial Differential Equations (PDEs) which preserve structure and accuracy under adaptation to unseen geometries. To this end, we introduce General-Geometry Neural Whitney Forms (Geo-NeWF): a data-driven finite element method. We jointly learn a differential operator and compatible reduced finite element spaces defined on the underlying geometry. The resulting model is solved to generate predictions, while exactly preserving physical conservation laws through Finite Element Exterior Calculus. Geometry enters the model as a discretized mesh both through a transformer-based encoding and as the basis for the learned finite element spaces. This explicitly connects the underlying geometry and imposed boundary conditions to the solution, providing a powerful inductive bias for learning neural PDEs which we demonstrate improves generalization to unseen domains. We provide a novel parameterization of the constitutive model ensuring existence and uniqueness of the solution. Our approach demonstrates state-of-the-art performance on several steady-state PDE benchmarks and provides a significant improvement over conventional baselines on out-of-distribution geometries.

Grokking, a sudden transition from memorization to generalization, has been closely linked to the emergence of low-dimensional representations; yet the mechanism driving this organization remains elusive. Here, we propose that intrinsic task symmetries are the key drivers of grokking, inducing structured geometries in representation space. Our analysis reveals a consistent three-stage training dynamic: (i) data memorization, (ii) intrinsic symmetry acquisition, and (iii) geometric organization. We show that generalization emerges during the symmetry acquisition phase, and subsequently the embedding space organizes into a low-dimensional structured geometry. We validate this mechanism across diverse algorithmic domains, spanning algebraic (modular arithmetic), structural (graph metric completion), and relational (comparison) reasoning tasks. Leveraging these insights, we formulate a symmetry-based criterion for generalization and propose symmetry- and geometry-prompting training strategies that can accelerate generalization. Together, our results establish intrinsic symmetry as a central mechanism enabling neural networks to move beyond memorization and achieve robust algorithmic reasoning.

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.

Motivated by recent work on benign overfitting in overparameterized machine learning, we study the generalization behavior of functions in Sobolev spaces $W^{k, p}(\mathbb{R}^d)$ that perfectly fit a noisy training data set. Under assumptions of label noise and sufficient regularity in the data distribution, we show that approximately norm-minimizing interpolators, which are canonical solutions selected by smoothness bias, exhibit harmful overfitting: even as the training sample size $n \to \infty$, the generalization error remains bounded below by a positive constant with high probability. Our results hold for arbitrary values of $p \in [1, \infty)$, in contrast to prior results studying the Hilbert space case ($p = 2$) using kernel methods. Our proof uses a geometric argument which identifies harmful neighborhoods of the training data using Sobolev inequalities.

Discrete diffusion models have recently emerged as a compelling alternative for language generation, enabling efficient non-autoregressive sampling while achieving strong empirical performance. A key design choice in discrete diffusion---absent in most continuous diffusion formulations---is the forward corruption kernel, with masked/absorbing corruption now dominating practice. Despite this empirical preference, there is limited statistical theory explaining when and why masking should outperform alternative kernels such as uniform replacement. In this paper, we take a step toward closing this gap from a statistical learning perspective. Our analysis establishes generalization bounds and, through an explicit comparison across different forward corruption kernels, reveals a central advantage of masking: it scales with the effective data support rather than the full ambient state space, thereby mitigating the curse of state space cardinality. We further derive structure-aware refinements that capture how concentration and sparsity in real sequential data sharpen the sample complexities. Together, these results offer a principled explanation for the empirical strength of masked diffusion and provide guidance for forward-kernel design in discrete generative modeling.

Model merging offers an efficient solution for integrating task-specific knowledge from multiple fine-tuned models. Most existing approaches focus on manipulating the difference vectors between fine-tuned and pretrained weights, often overlooking the generalization capabilities inherent in the pretrained parameters. In this work, we revisit the role of pretrained weights in model merging and investigate their efficacy from a subspace perspective. We find that the components of pretrained weights residing in the core subspace—defined by the dominant singular vectors—are essential for maintaining generalization across diverse tasks. Specifically, we present empirical evidence that pretrained weights are nearly first-order stationary and exhibit predominantly non-negative curvature within this core subspace with respect to multi-task loss landscapes, indicating near-optimality. These findings suggest that task-specific adaptations should be injected primarily into the orthogonal complement of the core subspace, thereby preserving the generalization properties of the pretrained model. Extensive experiments on vision and vision-language tasks show that this subspace-aware strategy consistently yields improvements over state-of-the-art training-free merging methods, including Task Arithmetic, LOT Merging, ISO, and TSV. The source code is available at https://github.com/SunWenJu123/model-merging.

Language models are pretrained on sequences that blend statistical regularities, which make text fluent, with factual associations between specific tokens, which encode knowledge of facts. While recent works suggest that generalization depends critically on the interaction of these two streams, such as the diversity of the contexts in which facts appear, these effects remain difficult to study systematically. This paper introduces a flexible synthetic testbed that combines a statistical stream of generic tokens with an abstract factual stream of source-target token pairs, enabling fine-grained control over their interaction, such as their composition into a sequence (contextual structure) or the level of context diversity carrying the facts at training time. Through controlled experiments, we find that higher contextual diversity delays in-distribution factual learning, while low diversity can harm out-of-distribution generalization in ways that depend on the contextual structure. As a result, the optimal diversity level depends on the training budget. Beyond factual recall failures, we also identify failures in statistical generalization as we study how the interplay between contextual design and diversity level impacts different aspects of generalization. Furthermore, through a series of controlled interventions on the model components, we trace failure in different aspects of generalization to distinct optimization bottlenecks, highlighting the importance of the embedding and unembedding layers. Overall, our synthetic framework allows us to isolate effects that would be confounded in large-scale studies, offering a controlled testbed for future investigations.

Data-driven algorithm design automates hyperparameter tuning, but its statistical foundations remain limited because model performance can depend on hyperparameters in implicit and highly non-smooth ways. Existing guarantees focus on the simple case of a one-dimensional (scalar) hyperparameter. This leaves the practically important, multi-dimensional hyperparameter tuning setting unresolved. We address this open question by establishing the first general framework for establishing generalization guarantees for tuning multi-dimensional hyperparameters in data-driven settings. Our approach strengthens the generalization guarantee framework for semi-algebraic function classes by exploiting tools from real algebraic geometry, yielding sharper, more broadly applicable guarantees. We then extend the analysis to hyperparameter tuning using the validation loss under minimal assumptions, and derive improved bounds when additional structure is available. Finally, we demonstrate the scope of the framework with new learnability results, including data-driven weighted group lasso and weighted fused lasso.

Deploying multimodal models in real-world scenarios requires generalization to new environments where recording conditions differ from training, a challenge known as multimodal domain generalization (MMDG). Standard architectures employ separate encoders for each modality and a fusion module, training the system end-to-end by optimizing on the fused features. In this paper, we identify that such joint optimization causes encoders to exploit cross-modal co-occurrences, statistical relationships between modalities that arise from source-specific recording conditions, rather than learning domain-invariant features. We term this failure mode Fusion Overfitting. To address this, we propose Modality-Entropy Regularization for Domain Generalization (MER-DG), which maximizes the entropy of each encoder's feature distribution to preserve feature diversity. MER-DG is architecture-agnostic and integrates into existing multimodal frameworks as an additive loss term. Extensive experiments on EPIC-Kitchens and HAC benchmarks demonstrate average improvements of ${\sim}5\%$ over standard fusion and ${\sim}2\%$ over state-of-the-art methods.

In recent years, there has been growing interest in understanding neural architectures' ability to learn to execute discrete algorithms, a line of work often referred to as neural algorithmic reasoning. The goal is to integrate algorithmic reasoning capabilities into larger neural pipelines. Many such architectures are based on (message-passing) graph neural networks (MPNNs), owing to their permutation equivariance and ability to deal with sparsity and variable-sized inputs. However, much existing work is either largely empirical and lacks formal guarantees or it focuses solely on expressivity, leaving open the question of when and how such architectures generalize beyond a finite training set. In this work, we propose a general theoretical framework that characterizes sufficient conditions under which MPNNs can learn an algorithm from a training set of small instances and provably approximate its behavior on inputs of arbitrary size with worst-case guarantees. Our framework applies to a broad class of algorithms, including single-source shortest paths, minimum spanning trees, and general dynamic programming problems, such as the $0$-$1$ knapsack problem. In addition, we establish impossibility results for a wide range of algorithmic tasks, showing that standard MPNNs cannot learn them and derive more expressive MPNN-like architectures that overcome these limitations. Finally, we refine our analysis for the Bellman–Ford algorithm, yielding substantially smaller required training sets and significantly extending the recent work of Nerem et al., 2025 by allowing for a differentiable regularization loss. Empirical results largely support our theoretical findings.

The pursuit of reducing the memory footprint of the self-attention mechanism in multi-headed self attention (MHA) spawned a rich portfolio of methods, e.g., group-query attention (GQA) and multi-head latent attention (MLA). The methods leverage specialized low-rank factorizations across embedding dimensions or attention heads. From the point of view of classical low-rank approximation, these methods are unconventional and raise questions of which objects they really approximate and how to interpret the low-rank behavior of the resulting representations. To answer these questions, this work proposes a generalized view on the weight objects in the self-attention layer and a factorization strategy, which allows us to construct a parameter efficient scheme, called Tucker Attention. Tucker Attention requires an order of magnitude fewer parameters for comparable validation metrics, compared to GQA and MLA, as evaluated in LLM and ViT test cases. Additionally, Tucker Attention~encompasses GQA, MLA, MHA as special cases and is fully compatible with flash-attention and rotary position embeddings (RoPE). This generalization strategy yields insights of the actual ranks achieved by MHA, GQA, and MLA, and further enables simplifications for MLA.

Continual learning (CL) aims to train models sequentially over multiple domains without forgetting previously learned knowledge. However, existing CL methods optimize for in-domain performance and are therefore prone to learning spurious, domain-specific cues ("shortcut learning"), which limits generalization to unseen domains after deployment. In this paper, we address this limitation through *continual learning of domain-invariant representation*. We introduce a broad class of CL methods that sequentially learn representations capturing invariant structures across domains. Our methods are motivated by the observation that such invariant structures often preserve the underlying causal mechanisms, which can reduce the risk of overfitting to domain-specific cues and thus offer better out-of-domain generalization. Our proposed CL methods combine replay-based training with a tailored sequential invariance alignment to learn---and preserve---invariant structures over time. We evaluate our methods under a deployment-oriented protocol that measures performance on unseen target domains. Across six benchmark and real-world datasets spanning vision, medicine, manufacturing, and ecology, our methods consistently outperform existing CL baselines in terms of generalization to unseen target domains. As an ablation, we further show that naïve extensions of sequential training with existing domain-invariant representation learning (DIRL) methods provide only limited benefits. To the best of our knowledge, this is the first work to develop domain-invariant representation methods for CL.

Imbalanced Domain Generalization (IDG) focuses on mitigating both *domain and label shifts*, both of which fundamentally shape the model's decision boundaries, particularly under heterogeneous long-tailed distributions across domains. Despite its practical significance, it remains underexplored, primarily due to the *technical* complexity of handling their entanglement and the paucity of *theoretical* foundations. In this paper, we begin by *theoretically* establishing the generalization bound for IDG, highlighting the role of posterior discrepancy and decision margin. This bound motivates us to focus on directly steering decision boundaries, marking a clear departure from existing methods. Then, we *technically* propose a novel Negative-Dominant Contrastive Learning (NDCL) for IDG to enhance discriminability while enforce posterior consistency across domains. Specifically, inter-class decision-boundary separation is enhanced by placing greater emphasis on negatives as the primary signal in our contrastive learning, naturally amplifying gradient signals for minority classes to avoid the decision boundary being biased toward majority classes. Intra-class compactness is encouraged through a re-weighted cross-entropy strategy, and posterior consistency across domains is enforced through a prediction-central alignment strategy. Finally, rigorous yet challenging experiments on benchmarks validate the effectiveness of our NDCL. The code is available at https://github.com/Alrash/NDCL.

This paper studies the problem of single domain generalization in time series classification, which aims to learn a generalized time series classification model using a single source domain. This problem is highly challenging due to unreliable supervision from domain scarcity. Although current approaches employ generative models for data augmentation, these synthesized samples often suffer from low diversity and intrinsic noise, leading to weak generalization ability. Towards this end, we propose a novel approach named Context-driven Diffusion with Progressive Expansion (CURE) for single domain generalization in time series classification. The core of our CURE is to generate semantic-aware and semantic-free contexts that jointly guide a conditional diffusion model for informative data expansion. In particular, our CURE first conducts representation disentanglement to extract semantic-aware and semantic-free representations from source data. To enhance generalizability through data synthesis, we not only retrieve reference time series trajectories with similar semantics for semantic-aware contexts, but also utilize adversarial strategies to learn semantic-free contexts. These contexts are integrated as joint conditions for a diffusion model, enabling diverse and reliable virtual data. To enhance expansion adaptability and stable optimization, we progressively update our semantic-free contexts via a memory bank and measure boundary properties for dynamic data filtering. Comprehensive experiments on benchmark datasets validate the effectiveness of CURE.

Recent work has found that neural networks with stronger generalization tend to exhibit higher representational alignment with one another across architectures and training paradigms. In this work, we show that models with stronger generalization also align more strongly with human neural activity. Moreover, generalization performance, model--model alignment, and model--brain alignment are all significantly correlated with each other. We further show that these relationships can be explained by a single geometric property of learned representations: the local intrinsic dimension of embeddings. Lower local dimension is consistently associated with stronger model--model alignment, stronger model--brain alignment, and better generalization, whereas global dimension measures fail to capture these effects. Finally, we find that increasing model capacity and training data scale systematically reduces local intrinsic dimension, providing a geometric account of the benefits of scaling. Together, our results identify local intrinsic dimension as a unifying descriptor of representational convergence in artificial and biological systems.

Feature learning strength (FLS), i.e., the inverse of the effective output scaling of a model, plays a critical role in shaping the optimization dynamics of neural nets. While its impact has been extensively studied under the asymptotic regimes---both in training time and FLS---existing theory offers limited insight into how FLS affects generalization in practical settings, such as when training is stopped upon reaching a target training risk. In this work, we investigate the impact of FLS on generalization in deep networks under such practical conditions. Through empirical studies, we first uncover the emergence of an *optimal FLS*---neither too small nor too large---that yields substantial generalization gains. This finding runs counter to the prevailing intuition that stronger feature learning universally improves generalization. To explain this phenomenon, we develop a theoretical analysis of gradient flow dynamics in two-layer ReLU nets trained with logistic loss, where FLS is controlled via initialization scale. Our main theoretical result establishes the existence of an optimal FLS arising from a trade-off between two competing effects: An excessively large FLS induces an *over-alignment* phenomenon that degrades generalization, while an overly small FLS leads to *over-fitting*.

Compositional generalization, the ability to recognize familiar parts in novel contexts, is a defining property of intelligent systems, yet modern models, despite massive training sets, see only a tiny fraction of the combinatorial input space. We ask what structure representations {must} have to support generalization to unseen combinations. We formalize three desiderata (divisibility, transferability, stability) and show they impose necessary geometric constraints under standard training: representations must decompose linearly into per-concept components, orthogonal across concepts. This grounds the Linear Representation Hypothesis as a necessary consequence of compositional generalization, and yields dimension bounds linking the number of composable concepts to embedding geometry. Empirically, across CLIP, SigLIP, and DINO, we find partial linear factorization with low-rank near-orthogonal per-concept factors, and the degree of this structure correlates with compositional generalization on unseen combinations. As models continue to scale, these conditions predict the geometry they may converge to. Code: https://github.com/oshapio/necessary-compositionality