NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

WavCIL: Wavelet Coefficient-Domain Invariant Learning for Dynamic Graph OOD Generalization

Yantong Zhu, Weining Shi, Keyi Li, Xiaoyan Xie, Yi Chang, Qinggang Zhang, Tianjiao Pu, Ji Qiao, Zhihong Zhang

Expand abstract

Dynamic Graph Neural Networks (DyGNNs) demonstrate powerful representation capabilities in crucial time-series systems by leveraging graph structure and temporal dynamics. However, existing dynamic graph models often suffer from unpredictable distribution shifts induced by complex time-varying factors such as emergency incidents, leading to severe performance degradation due to their low out-of-distribution (OOD) generalization. Graph invariant learning has been extensively studied to solve this problem through feature disentanglement in the time or spectral domain. Despite recent advances, existing methods suffer from critical limitations in practice: Temporal-based methods struggle to disentangle transient perturbations and stable patterns that are highly overlapping and entangled within the same period. While spectral-based methods attempt spectral-domain disentanglement via the global Fourier transform, these methods inherently lose temporal localization and fail to effectively characterize distribution shifts driven by time-varying factors. To this end, in this paper, we propose **Wav**elet **C**oefficient-Domain **I**nvariant **L**earning **(WavCIL)**, a novel framework for resolving distribution shifts by conducting disentanglement within a wavelet coefficient domain endowed with time-frequency localization capabilities. Specifically, WavCIL consists of two key components: (i) a wavelet transform module, which leverages a set of wavelet bases to map the input temporal signals into wavelet coefficient domain, simultaneously capturing frequency components and localizing their temporal occurrences; and (ii) a coefficient domain invariant learning module, which utilizes an adaptive mask mechanism to disentangle stable causal patterns invariant to time-varying factors from spurious patterns driven by time-varying factors in the coefficient domain, guiding the prediction process to rely more on the causal patterns. Extensive experiments on several benchmark datasets demonstrate that WavCIL achieves state-of-the-art generalization performance when handling distribution shifts. Our code is released at https://anonymous.4open.science/r/WavCIL-4B8C.

Poster
Verifiable Environments Are LEGO Bricks: Recursive Composition for Reasoning Generalization

Hao Xiang, Qiaoyu Tang, Le Yu, Yaojie Lu, Xianpei Han, Ben He, Le Sun, Bowen Yu, Peng Wang, Hongyu Lin, Dayiheng Liu

Expand abstract

Reinforcement Learning (RL) with verifiable environments has emerged as a powerful approach for enhancing the reasoning capabilities of Large Language Models (LLMs). While prior research indicates that scaling environment quantity improves RL performance, existing manual or individual construction methods suffer from linear scaling limits, thereby hindering scalable reasoning generalization. This paper introduces RACES (\textbf{R}ecursive \textbf{A}utomated \textbf{C}omposition for \textbf{E}nvironment \textbf{S}caling), a framework that treats verifiable environments as composable building blocks that can be recursively assembled. The key insight is that when the codomain (output type) of one environment matches the domain (input type) of another, they can be automatically fused into a new verifiable environment, enabling recursive composition. RACES is instantiated with 300 individual environments and defines a set of composition operators (Sequential, Parallel, Sort, and Select) that induce diverse reasoning patterns. Extensive experiments show that RL training on these composite environments consistently enhances reasoning generalization. Specifically, RACES improves DeepSeek-R1-Distill-Qwen-14B by an average of 3.1 points (from 48.2 to 51.3) and boosts Qwen3-14B performance from 58.8 to 61.1 on five benchmarks, which are unseen during the construction of training environments. Moreover, RACES achieves performance comparable to training on 300 individual environments using only 50 base environments, demonstrating significant efficiency in environment utilization.

Poster
Unlocking Compositional Generalization in Continual Few-Shot Learning

Phu Quy Nguyen Lam, Phu Hoa Pham, Dao S Minh, Chi Nguyen Tran, Trung Kiet Huynh, Long Tran-thanh

Expand abstract

Object-centric representations promise a key property for few-shot learning: Rather than treating a scene as a single unit, a model can decompose it into individual object-level parts that can be matched and compared across different concepts. In practice, this potential is rarely realized. Continual learners either collapse scenes into global embeddings, or train with part-level matching objectives that tie representations too closely to seen patterns, leaving them unable to generalize to truly novel concepts. In this paper, we identify this fundamental structural conflict and pioneer a new paradigm that strictly decouples representation learning from compositional inference. Leveraging the inherent patch-level semantic geometry of self-supervised Vision Transformers (ViTs), our framework employs a dual-phase strategy. During training, slot representations are optimized entirely toward holistic class identity, preserving highly generalizable, object-level geometries. At inference, preserved slots are dynamically composed to match novel scenes. We demonstrate that this paradigm offers dual structural benefits: The frozen backbone naturally prevents representation drift, while our lightweight, holistic optimization preserves the features' capacity for novel-concept transfer. Extensive experiments validate this approach, achieving state-of-the-art unseen-concept generalization and minimal forgetting across standard continual learning benchmarks.

Poster
Understanding Generalization Requires Universal Induction

Aram Ebtekar, Marcus Hutter, Danica J. Sutherland

Expand abstract

Classical statistical theory is insufficient to explain the successes of general-purpose AI models, because it depends on handcrafted inductive biases that it cannot justify. No Free Lunch theorems force any learner that beats chance on some environments to underperform on others, and meta-learning which environments are more likely only pushes the problem up a level. The inductive bias must therefore be grounded in something other than data. By deriving its bias from Turing universality, Solomonoff induction (SI) competes with all computable learners. However, its regret bounds include large ``constants,'' such as the size of a learner's (losslessly compressed) full codebase. We relativize SI to an information vantage point, which includes all pre-existing code and data. This reframes the inductive bias: instead of favoring some absolute notion of simplicity, we favor accessibility with respect to our vantage point. The relativized SI is sample-optimal on finite data: any learner that outperforms it necessarily contains inaccessible information about the data. While SI is incomputable and hence not a practical algorithm, it provides a formal optimum for inference in the limit of infinite compute. We argue that algorithmic information theory, which underlies SI, is necessary to explain the generalization behavior of modern (and future) AI systems.

Poster
Transformers Can Learn Multiclass Classification In-Context: Isotropy Governs Generalization

Daehan Yoon, Chulhee Yun

Expand abstract

In-context learning plays a central role in transformer-based large language models, yet its theoretical understanding remains limited. In this work, we study multiclass in-context classification under more realistic settings, including anisotropic class centers and label imbalance, by introducing a spectral data generation framework that constructs class-center matrices with a prescribed singular spectrum. We first show that the isotropy of class-center vectors, as quantified by the stable rank, improves ICL generalization performance. From a meta-learning perspective, our theorem shows that linear transformers can learn multiclass in-context classification with near-optimal per-label sample complexity, extending prior guarantees beyond the binary setting. In the test label imbalance regime, our analysis reveals that queries from majority classes are easier to classify, while those from minority classes are more error-prone; moreover, robustness to this bias improves with the stable rank. Finally, we empirically demonstrate that our theory is consistent with observations on transformers and pretrained large language models.

Poster
Transductive Generalization for GNNs via Optimal Transport

MoonJeong Park, Seungbeom Lee, Kyungmin Kim, Jaeseung Heo, Seunghyuk Cho, Shouheng Li, Sangdon Park, Dongwoo Kim

Expand abstract

Analyzing the generalization of graph neural networks (GNNs) on node classification is challenging: message passing induces dependencies that break the i.i.d. assumption underlying standard inductive generalization theory. While distribution-free transductive theory resolves this dependency issue, existing classical bounds are either intractable or correlate poorly with empirical generalization. To fill this gap, we propose a transductive generalization bound based on optimal transport, connecting generalization to intra-class concentration and inter-class separation in the feature space. Empirically, our bound aligns strongly with the generalization gap. Building on this representation-based approach, we analyze how feature geometry evolves under repeated message passing. This explains the non-monotonic relationship between GNN depth and generalization, providing new insights into oversmoothing research. Practically, our bound serves as a geometry-aware regularizer that corrects the harmful behavior of message passing, yielding consistent performance improvements.

Poster
Towards Better Generalization in Lifelong Person Re-Identification with Flatness-Aware Learning

Seungbin Hong, Sung Whan Yoon, Jae-Young Sim

Expand abstract

Lifelong person re-identification (LReID) requires models to continuously learn from sequentially arriving domains while retaining discriminative power for previously seen identities. A key challenge is to prevent catastrophic forgetting without access to old data, especially under exemplar-free constraints. While flat-minima optimization has shown promise in continual learning, we identify a gradient conflict that limits the effectiveness of standard Sharpness-Aware Minimization (SAM) in LReID. The ReID loss gradient dominates the perturbation direction, causing the sharpness of the distillation loss to be underestimated, which hinders the flattening of the landscape necessary for knowledge retention. To resolve this conflict, we propose a framework that unifies selective flatness-aware optimization, dual-model training, and weight-space interpolation. Specifically, we maintain two models per task: a stability model whose SAM perturbation is computed solely from the distillation loss, and a plasticity model optimized for the current domain. By decoupling the perturbation objectives, our selective SAM achieves more targeted sharpness exploration along the distillation landscape, guiding the stability model toward flatter and robust regions. After training, the two models are fused via weight-space interpolation, and we provide a theoretical bound showing that flatter stability solutions tighten the interpolation bound and reduce forgetting during fusion. Our method is lightweight, modular, and readily compatible with existing LReID frameworks. Experimental results demonstrate that the proposed method consistently improves the performance of LReID in terms of both knowledge retention and generalization to unseen domains.

Poster
Tokenizer Choice Shapes Generalization in State-Centric Learning for Planning

Vishal Pallagani, Nitin Gupta, John A Aydin, Biplav Srivastava

Expand abstract

Planning has traditionally been addressed with symbolic search guided by domain models and heuristics. Learning offers a complementary way to exploit structure shared across related problems rather than solving each instance from scratch, making generalization to unseen instances a central challenge for learned planners. Recent work has explored planning mainly through action-centric models that predict actions or full plans from problem descriptions. A parallel line of work instead learns goal-conditioned transition models over states (state-centric learning) and recovers actions by matching predicted successor states to symbolic successors. In this setting, however, prior work has relied mainly on Weisfeiler-Leman (WL) representations, leaving the role of tokenizer choice unclear. We present a controlled study of tokenizer choice in state-centric learned planning, comparing WL, shortest-path, GraphBPE, SimHash, and a deterministic baseline within a common pipeline. Across varied benchmark domains and multiple planner configurations, WL performs well overall, but no tokenizer is best everywhere: the leading representation changes by domain, and tokenizer differences are larger on extrapolation rather than on interpolation tasks. Across the studied configurations, tokenizer choice is the dominant factor in held-out generalization: it explains 57% of performance variance, more than predictor architecture or any other pipeline factor, and its effect is substantially amplified on out-of-distribution problems relative to in-distribution ones. We conclude that representation choice is a strong determinant of generalization in state-centric learned planning and that tokenizer choice is a first-order modeling decision rather than a fixed preprocessing step. Code: https://tinyurl.com/4hre257u

Poster
TF-PRVR: Training-Free Partially Relevant Video Retrieval for Real-World Generalization

Giyeol Kim, Chanho Eom

Expand abstract

Partially Relevant Video Retrieval (PRVR) aims to retrieve untrimmed videos containing moments relevant to a given text query. Despite recent progress, existing PRVR methods suffer from two key limitations: a fixed video decomposition scheme that causes semantic dilution, and weak cross-domain robustness due to task-specific training. In this paper, we propose TF-PRVR, the first training-free framework for PRVR designed to improve real-world generalization. TF-PRVR leverages frozen vision-language features to construct video-specific hierarchical representations. It derives temporal semantic signals from frame-level features and applies frequency-based multi-scale analysis to identify adaptive temporal boundaries, producing hierarchical segments with coherent event-level semantics. Built on these segments, TF-PRVR constructs a unified multi-scale graph and propagates query relevance across temporally and semantically related nodes. A moment-aware scoring strategy then aggregates temporally aligned relevance across scales, emphasizing consistently supported moments while suppressing isolated false responses. Without task-specific training, TF-PRVR preserves the general-purpose alignment capability of pre-trained vision-language models and avoids dataset-specific overfitting. Extensive experiments on standard PRVR benchmarks demonstrate competitive retrieval performance and strong robustness under cross-domain evaluation, suggesting a practical direction for real-world PRVR.

Poster
StomataBench: Measuring Taxonomic Generalization in Stomatal Detection

Sainath Reddy Gummi, Hankui K. Zhang, Chulwoo Pack, Shyam Solanki, Young Chang

Expand abstract

Accelerating stomata measurement is essential for developing drought-resilient plants. However, most automated stomata detectors are tested under narrow imaging and species conditions. Their behavior under broader ecological distribution shift remains unclear. We introduce StomataBench, a benchmark for stomatal object detection across species, acquisition protocols, and biological density regimes. The benchmark uses a curated multispecies training set of 3,154 images with 134,802 annotated stomata from 21 woody species. It is evaluated on four test datasets totaling 19,657 images, covering diverse woody taxa and heterogeneous public crop sources. Unlike standard detection benchmarks that report only AP, StomataBench evaluates measurement reliability. We measure COCO AP, F1, score-threshold sensitivity, image-level count error, and stomatal density agreement. We benchmark sixteen detectors spanning two-stage, one-stage, transformer, open-vocabulary, and promptable foundation models. The results show that model rankings change with the biological objective. GDINO-SB provides the strongest average AP generalization, while SAM3-OD gives reliable operating point for count and density estimation. Conventional detectors remain competitive on near-domain woody plants, but foundation models are robust under harder ecological shifts. Failure analysis shows that far-OOD detection is often limited by missing stomata. These results identify cross-lineage proposal generation as a key unsolved problem for reliable stomatal phenotyping under ecological distribution shift.

Spotlight
State-Action Selection in General Stochastic Games under Regularized Policy Gradient Methods

Omar Abbadi, Rida Laraki, Panayotis Mertikopoulos

Expand abstract

We study a general class of regularized policy gradient methods for episodic learning in general-sum stochastic games, with the aim of describing the long-run behavior of the induced policy iterates. In this challenging setting, learning even the simplest pointwise solution concepts is known to be intractable. This leads us to switch our focus to less restrictive *setwise* solution concepts: rather than only asking whether the iterates converge to a specific policy, we ask which *sets of policies* are selected by the learning dynamics in the long-run. In one-shot games, this setwise perspective admits a crisp characterization: a product set of pure strategies is asymptotically stable if and only if it is *closed under better replies* (club). We show that the direct statewise analogue of this principle *fails* outright in stochastic games because actions affect not only current rewards but also the future evolution of the game. To restore this principle in a transition-aware form, we introduce $Q$-*club* sets, namely state-action sets in which every *allowed* action has a strict $Q$-value advantage over every *excluded* action, thereby providing a setwise generalization of *strict Nash policies*. We then prove that $Q$-club sets are stochastically stable and attracting under the learning dynamics, even in the presence of noise and uncertainty. Finally, we establish local convergence rates to $Q$-club sets governed by the regularizer; in particular, entropic regularization yields exponential convergence, while Euclidean regularization yields finite-time convergence.

Poster
Skipping Domain Shifts: Domain Memory Retention Enhanced Hyperspectral Single-Source Domain Generalization

Taiqin Chen, Hao Sha, Xiaochen Feng, Chunyu Chen, Xiangyu Chen, Ke Chen, Yongbing Zhang

Expand abstract

Hyperspectral single-source domain generalization aims to train a robust model capable of overcoming domain shifts using only one available source domain for cross-scene hyperspectral image (HSI) analysis. Existing methods typically adopt data augmentation to broaden decision boundaries or utilize style transfer techniques for test-time alignment. However, these studies either fail to completely bridge the domain gap or inadvertently disrupt the discriminative spectral information, leading to suboptimal generalization and degraded classification performance on the target domain. To overcome these limitations, we propose a novel method, termed the domain memory retention framework, which optimizes domain memory during training to explicitly encode the semantic structure of the source domain. During inference, the target features are projected to the source distribution by leveraging this memory, thereby skipping domain shifts and improving generalization performance across unseen scenes. Specifically, the memory is formulated as a codebook, consisting of multiple semantic codewords. We further develop a semantic richness constraint and a perturbation robustness constraint to enhance representation diversity while empowering domain invariance for the memory. Extensive experiments conducted on three remote sensing datasets demonstrate that the proposed method outperforms state-of-the-art methods.

Poster
Rethinking Memorization–Generalization Trade-Off in Generative Models

Jiseok Chae, Kyuwon Kim, Donghwan Kim

Expand abstract

Standard generative models face a memorization--generalization trade-off, that is avoiding memorization is considered necessary for generalization. In supervised learning, however, recent studies show that highly overparametrized models can memorize (interpolate) training data while still generalizing well, a phenomenon known as benign overfitting. Motivated by this, we investigate whether generative models can similarly bypass the trade-off and generalize while interpolating. Generative models should minimize the distance between the true data distribution and the distribution induced by mapping the full latent distribution through the generator. But since the true distribution is inaccessible, existing models instead minimize empirical risk with respect to the training distribution. Under this standard formulation, exact empirical-risk minimization forces the generator to produce only training samples. To address this issue, we consider an alternative empirical risk based on presampled latent variables. Our experiments demonstrate that the resulting presampled-latent version of the flow matching model exhibits benign overfitting on standard image benchmarks, such as MNIST and CIFAR-10. As a theoretical proof of concept, we recast generative modeling as a regression problem and extend existing benign overfitting theory to our setting. Together, these results establish, for the first time, that benign overfitting can occur in generative models.

Poster
Random-Effects Centroids for Domain Generalization

Jinyong Jeong, Hyungu Kahng, Seoung Bum Kim

Expand abstract

Domain generalization methods train on multiple source environments but usually deploy a single pooled predictor. We study a different deployment object: one head per source environment on a shared representation, combined by fixed weights when the test environment label is unavailable. Under a same-meta random-effects model, where environment-specific population minimizers are iid deviations around a shared meta-mean, the optimal fixed simplex rule is the uniform centroid. The same analysis separates this centroid from pooled ERM: under a common feature second moment, pooled ERM is sample-weighted, while the centroid is environment-weighted, giving a closed-form gap controlled by sample-weight imbalance. We introduce **REC** (Random-Effects Centroid), a training objective that fits source-specific experts while directly scoring their averaged deployment head. In the common-second-moment case, the centroid term controls the trained-centroid discrepancy and the dispersion of environment-specific optima. Synthetic experiments isolate the population estimator effect and learned-representation behavior; WILDS experiments show that REC is competitive with established DG baselines and that the centroid term is necessary for a useful averaged predictor.

Poster
PENEX: AdaBoost-Inspired Neural Network Regularization

Klaus-Rudolf Kladny, Bernhard Scholkopf, Michael Muehlebach

Expand abstract

AdaBoost sequentially fits so-called weak learners to minimize an exponential loss, which penalizes misclassified data points more severely than other loss functions like cross-entropy. Paradoxically, AdaBoost generalizes well in practice as the number of weak learners grows. In the present work, we introduce Penalized Exponential Loss (PENEX), a new formulation of the multi-class exponential loss that is theoretically grounded and, in contrast to the existing formulation, amenable to optimization via first-order methods, making it a practical objective for training neural networks. We demonstrate that PENEX effectively increases margins of data points, which can be translated into a generalization bound. Empirically, across computer vision and language tasks, PENEX improves neural network generalization in low-data regimes, matching and in some settings outperforming established regularizers at comparable computational cost. Our results highlight the potential of the exponential loss beyond its application in AdaBoost.

Poster
Orthogonal Origin Parking: Decoupling Lorentz Manifolds for Robust OOD Generalization

Peter J Kampen, Anders N. Christensen, Morten Rieger Hannemose, Anders Dahl, Josefine Vilsbøll Sundgaard

Expand abstract

Fine-grained classification in deep taxonomies suffers from representational crowding: large macro-classes dominate the feature space, leaving little angular capacity for fine-grained distinctions, particularly under out-of-distribution (OOD) shifts. Standard hyperbolic spaces offer exponential volume for tree-like data, but embedding an entire taxonomy into a single manifold forces distinct macro-branches to compete for shared angular capacity and entangles their optimization updates. We introduce a geometric architecture that embeds hierarchical taxonomies into a Cartesian product of Lorentz manifolds. Our core mechanism, \textit{Orthogonal Origin Parking} (OOPark), decouples the taxonomy at the root level by assigning an independent Lorentz manifold to each macro-branch and penalizing inactive embeddings that deviate from the manifold origin. We evaluate across four biological OOD tasks/benchmarks (skin lesions, iWildCam, fungi, and plankton) characterised by hierarchical class structure, class imbalance, and domain shift. OOPark generally outperforms standard Euclidean and single-manifold hyperbolic baselines while using 20-dimensional sub-manifolds against Euclidean baselines of up to 512 dimensions, preserving both fine-grained accuracy and global taxonomic fidelity under severe distribution shifts. Code is provided in the supplementary material; a public release will follow upon acceptance.

Poster
Nexus: Same Pretraining Loss, Better Downstream Generalization via Common Minima

Huanran Chen, Huaqing Zhang, Xiao Li, Yinpeng Dong, Ke Shen, Jun Zhu

Expand abstract

The foundational capabilities of large language models are acquired during pretraining on internet-scale, highly heterogeneous data mixtures. In this work, we investigate a geometric question regarding the converged state of this pretraining process: Does the model converge to a common minimizer across all data sources (e.g., \cref{fig:cwa_illustration:close}), or merely a minimizer of the averaged loss (e.g., \cref{fig:cwa_illustration:distant})? We hypothesize that the geometric ``closeness'' of task-specific minima is intrinsically linked to downstream generalization. However, we reveal that standard optimizers (e.g., AdamW) often converge to points where task-specific minima are distant from each other. To address this, we propose the \textit{Nexus} optimizer, which encourages the closeness of these minima by maximizing gradient similarity during optimization. Extensive experiments across models ranging from 130M to 3B parameters demonstrate that Nexus significantly boosts downstream performance, despite achieving nearly \textit{the same pretraining loss} (see \cref{fig:demo:benchmark}). Notably, on the 3B model, Nexus reduces the out-of-distribution loss by 0.012 and yields up to a 15.0\% accuracy improvement on complex reasoning tasks (e.g., GSM8k). Our findings challenge the reliance on pretraining loss as the sole proxy for model evaluation and highlight the critical role of optimizer implicit bias in unlocking downstream generalization.

Poster
Neural population tuning statistics as priors for multitask generalization

Zach Cohen, Jan Drugowitsch

Expand abstract

Neural populations exhibit heterogeneous tuning: Neurons differ in their selectivity, gain, and in how their preferred inputs are distributed across stimulus or latent-variable space. While this heterogeneity has been studied extensively in the context of encoding precision, its implications for learning and generalization are less well-understood. Here, we develop a probabilistic theory that links a neural population's tuning statistics to its inductive bias: in large populations, the tuning statistics induce a kernel and, equivalently, a prior of a generative model over functions. This prior in turn predicts the population's ability to generalize across a distribution of related tasks. Our theory determines the population tuning statistics that minimize the generalization error over a distribution of functions. We further present a normative account of tuning adaptation inspired by our probabilistic approach, showing that tuning changes that maximize the marginal likelihood of observed data systematically reduce generalization error relative to non-adapting populations. Applied to the adaptation of hippocampal spatial representations in a reward learning task, we found that our theory could explain the over-representation of rewarding locations through experience in an environment with a localized reward. We further demonstrate how different components of the population's generative model differentially impact the speed and magnitude of tuning adaptation. Overall, our work proposes a novel connection between tuning variability and a population's ability to learn and to generalize across tasks and develops a Bayesian theory of tuning adaptation in neural populations during learning.

Oral
Multiple Descent of Generalization Curve for Optimally Regularized Ridge Regression

Maxim Bochkov, Fedor Noskov

Expand abstract

It is known that optimally regularized ridge regression can exhibit non-monotone dependence of the generalization risk on the sample size. We show that this behavior can be arbitrarily intricate in finite dimensions: for every $d \ge 2$, there exists a random design linear model in $\mathbb{R}^d$ whose generalization risk, viewed as a function of the sample size, exhibits at least $d - 1$ local maxima.

Poster
Learning Optimal Transport Plans Via Autoregressive Token Regression

Ivan Jaen Marquez, Takis Chytas, Vikas Singh

Expand abstract

Optimal Transport (OT) provides a principled framework for comparing probability distributions. Classical solvers remain the gold standard for this task, providing exact solutions. At the same time, OT is an ideal testbed for studying whether learning-based amortized models can internalize the structure of OT type problems. Here, ground-truth solutions are available, constraints are explicit, and exact discrete solutions have well-characterized sparsity. We describe ToR-OT, a transformer-based model that learns to predict discrete OT plans through autoregressive token regression. We reframe OT into a structured sequence generation task. ToR-OT decodes sparse transport plans as sequences rather than solving an optimization task at test time. We focus on three properties: (a) Zero-shot generalization: pretraining on a large synthetic corpus of OT instances allows a single model to generalize to unseen distribution pairs and heterogeneous cost metrics without retraining. (b) Sparse primal decoding: by exploiting the inherent sparsity in exact discrete OT solutions, the decoder predicts only nonzero entries (avoids dense coupling outputs). (c) Amortized inference: on low-dimensional discrete problems, ToR-OT gives accurate approximate plans and compares favorably with per-instance neural OT baselines. Overall, ToR-OT is not yet a replacement for classical solvers, but best viewed as a step toward understanding how autoregressive transformers can learn structure of constrained optimization problems like OT.

Poster