NeurIPS 2026 Papers

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.

Precipitation downscaling over complex terrain is essential for resolving water availability, flood risk, and extremes, but current evaluations rarely test whether models learn orographic structure rather than only pixel-wise accuracy. We introduce OroPrecipBench, a km-scale benchmark for orographic precipitation downscaling across five U.S.\ terrain regimes. The benchmark provides a shared-domain dual-track protocol: a perfect-model track that coarsens 3\,km HRRR analyses to isolate controlled super-resolution skill, and an observation-grounded track that downscales 25\,km ERA5 to 4\,km PRISM to test practical robustness under source-target shift. We also introduce PEP-D, a terrain-conditioned diagnostic that measures errors in elevation--precipitation profiles, complementing pointwise, spatial, and extreme-precipitation metrics. Evaluating ten statistical, deterministic deep-learning, and generative baselines shows that strong controlled super-resolution does not reliably translate to real-case or cross-regime downscaling. Models that reduce pixel-wise error can still distort terrain-dependent precipitation and upper-tail rainfall, and terrain-conditioned structure is especially fragile under transfer. OroPrecipBench exposes these failure modes and motivates source-shift-robust, terrain-constrained, and tail-aware downscaling methods. Benchmark is available \href{https://anonymous.4open.science/r/OroPrecipBench-612D}{here}.

We study the Sobolev IPM problem for probability measures supported on a graph metric space, where critic function is constrained to lie within the unit ball defined by Sobolev norm. Sobolev IPM is intrinsically coupled with $L^p$ geometric structure within its definition, limiting its ability to incorporate other prior geometry beyond the $L^p$ paradigm. Conversely, the classic optimal transport is flexible, easy to adapt to various geometric structures by simply changing its ground cost. An important example is Orlicz-Wasserstein (OW) which utilizes \emph{Orlicz} geometric structure to generalize the $L^p$ within the standard $p$-order Wasserstein, and remarkably play a vital role to advance machine learning methodologies. Inspired by recent advantages of OW, in this work, we leverage a specific class of convex functions for Orlicz geometry to mitigate such limitation for Sobolev IPM, and propose the generalized Sobolev IPM (GSI). Our GSI approach encompasses Sobolev IPM as a special case while accommodating diverse geometric priors beyond $L^p$. It however brings up significant computational hurdles that compound those already notoriously inherent in Sobolev IPM. To address these challenges, we theoretically establish a novel connection between \emph{Orlicz-Sobolev} norm and \emph{Musielak} norm which facilitates a novel efficient regularization for GSI. By further exploiting the underlying graph structure, we show that the regularized GSI reduces to a simple univariate optimization problem, achieving notably computational efficiency, enabling its usage in practical applications. We empirically illustrate that the regularized GSI is several-order faster than the popular OW in computation, and its performances compare favorably to other transport baselines for comparing graph-based measures in document classification and topological data analysis.

General-purpose embodied AI systems are increasingly expected to interpret multimodal instructions and perform fine-grained physical tasks. However, many manipulation benchmarks provide limited support for jointly evaluating procedural correctness, local geometric sensitivity, and human-guided correction. We propose origami as a benchmark based on real images for procedural state-transition reasoning in embodied AI. Origami requires precise operations such as tucking, inserting, reversing, and folding overlapping layers, making it a focused domain for studying physical state transitions. We introduce a dataset of step-by-step folding images, textual entries, and difficulty annotations, and define two tasks: autonomous state transitions from multimodal instructions and human-guided collaborative transitions. Unlike crease-pattern- or diagram-based origami benchmarks, our setting exposes agents to actual intermediate folding states. For an initial simulator-based evaluation, we use a multimodal LLM controller to evaluate procedural understanding, action-interface grounding, and retry-based correction.

I show that ordinary least squares (OLS) predictions can be rewritten as a restricted attention module, akin to those at the heart of transformer architectures. This connection emerges by reframing OLS as a similarity-based prediction rule operating in a learned embedding space. From this perspective, least squares does not estimate coefficients per se, but instead selects an embedding that minimizes squared prediction error by matching training and test vectors through inner products — mapping naturally onto the query-key-value structure of attention. I then extend the framework to dimensionality reduction, nonlinearity, and connections to time series econometrics. Monte Carlo simulations and real-data experiments on UCI/OpenML benchmarks show that a direct implementation of nonlinear Attention Regression performs competitively against standard machine learning baselines. In the other direction, within a transformer architecture for tabular data, an attention block can be replaced by an explicit regression on polynomial-expanded features, matching or exceeding the standard transformer's predictive accuracy at a fraction of its parameter count.

Goal-conditioned planning requires compressing value functions into low-dimensional representations, yet which property of the compression best predicts planning quality remains unclear. We show that reconstruction accuracy and ordinal fidelity---how faithfully the ranking of successor states by value is preserved---are two projections of the same approximation error: in smooth geometry they couple tightly and $L^2$ is a sufficient summary; when topology creates dense small-gap regions they decouple, and the ordinal projection carries complementary planning-relevant information beyond $L^2$. Through experiments on gridworld, continuous navigation (640 runs, 8 environments), and Maze2D (240 runs, 3 D4RL layouts), we establish three results. First, neighbor-restricted ordinal fidelity ($\tau^{\mathrm{nbr}}$) adds $\Delta R^2 = 0.106$ of incremental explanatory power beyond $L^2$ after environment fixed effects. Second, monotone recalibration degrades $L^2$ by up to $\sim\!400\times$ while preserving $\tau^{\mathrm{nbr}}$ and planning success. Third, in matched-pair comparisons controlling for $L^2$, higher $\tau^{\mathrm{nbr}}$ predicts better planning 76\% of the time. We provide theoretical grounding via a gap-weighted Bellman-excess decomposition that links reconstruction error, local ordinal inversions, and planning regret.

Proximal Policy Optimization (PPO) stabilizes policy updates by clipping the importance ratio between the new policy and the data-collecting policy. The standard rule uses one fixed clipping range for all samples. While this choice offers simplicity and effectiveness, it fails to distinguish statistically harmful extreme ratio values from large, yet tolerable changes. We study how to adapt the clipping radius to the reliability of the induced importance ratio distribution. Our case study demonstrates that reliable updates require controlling these extreme ratios, rather than merely verifying that individual ratios fall within a fixed interval. Motivated by this observation, we propose Ordered Policy Optimization (\textsc{OPO}). In each minibatch, \textsc{OPO} sorts the importance ratios, selects the extreme ratios that potentially have the largest effect on the update, and assigns them adaptive sample-wise radii instead of using the fixed clipping radius. These radii are derived from a prescribed tail profile, which specifies desired target values for the selected extreme ratios. The resulting objective remains first-order and requires only minibatch sorting and a modified policy loss. Experiments conducted on MuJoCo and Atari benchmarks demonstrate that this simple modification achieves strong performance relative to established baselines.

We propose a structure learning method for directed acyclic graph (DAG) models of zero-inflated multivariate count data. Although causal structure learning with zero-inflated Poisson models is identifiable, existing algorithms rely on local search over the DAG space, leading to poor scalability and a tendency to converge to suboptimal solutions. To address these limitations, we develop a novel order-based scoring framework that searches over variable orderings rather than directly over DAGs. Building on this framework, we further extend the method to multiple heterogeneous datasets by identifying a shared variable ordering while allowing each dataset to have its own edge set. We implement a stochastic hill climbing algorithm with a random-to-random (R2R) proposal operator to efficiently explore the ordering space. Simulation studies show that the proposed method outperforms competing approaches and improves graph estimation accuracy under increasing heterogeneity when a shared causal ordering is preserved. We apply the method to single-nucleus RNA sequencing data from a major depressive disorder study.

Flow-based generative models have emerged as an efficient alternative to diffusion models. We study the sample complexity of learning a target distribution in Wasserstein distance via flow matching, which learns a velocity field transporting a base distribution to the target. Under standard assumptions on the network architecture and data distribution, we show that any squared-loss flow-matching variant including rectified flow with bounded derivatives achieves $\widetilde{O}(\varepsilon^{-2})$ sample complexity with the Wasserstein guarantee taking the form $W_2 \le K(\varepsilon + C\sqrt{\varepsilon_{app}})$, where $\varepsilon_{app}$ is the approximation error of the network class. This improves upon existing $\widetilde{O}(\varepsilon^{-4})$ guarantees and does not require the Polyak \L{}ojasiewicz condition assumed in prior work. Instead, we exploit a structural property of the squared-loss objective - it induces an approximate Bernstein condition that directly ties the variance of the excess loss to the excess risk up to additive approximation-error. This condition enables a localized Rademacher complexity analysis yielding fast $O(1/n)$ rates. The flow-matching schedule enters only through the Lipschitz constant of the pointwise loss, which we characterize explicitly for standard schedules.

Masked Diffusion Models (MDMs) are trained under an objective that is symmetric over decoding orders, but at sampling time a single order must be committed to. A common task in such models is to evaluate the log-probability of a candidate sequence, either as a downstream score or to rank a finite pool of candidates. We observe that this evaluation depends substantially on the decoding order chosen, even though the training objective averages over orders. For both a model trained on closed-form Probabilistic Context-Free Grammar (PCFG) data and a pre-trained MDM on OpenWebText, the per-order log-probability rankings of fixed candidate pools disagree on most inputs. We characterize this phenomenon and propose Order-Marginalized Scoring (OMS), a simple Monte-Carlo (MC) estimator of the order-marginalized log-likelihood $\tilde \mu(x)=E_\sigma[\log p_\sigma(x)]$. We use it in two complementary algorithms: a reranker over a fixed candidate pool, and a stochastic best-of-$N$ decoder that combines adaptive position selection with $\tilde\mu$-reranking. On Sudoku, the decoder substantially raises exact-match accuracy over the strongest single-order baseline from $65.3$% to $94.7$%. On a zero-shot protein-stability benchmark, the reranker improves Spearman correlation between log-probability scores and experimental folding stabilities over a per-order baseline by $0.07$. These results suggest that OMS is a useful alternative to per-order scoring for order-agnostic problems.

DETR-style detectors use one-to-one bipartite matching during training to assign object queries to ground-truth objects, enabling end-to-end set prediction without non-maximum suppression (NMS). However, without an explicit de-duplication procedure, multiple queries can still produce highly similar hypotheses for the same object, making training unstable and predictions less decisive. Inspired by the sequential ordering of NMS, we propose DETRNN, a plug-and-play module that turns unordered object queries into a competition-aware sequence for recurrent refinement. DETRNN builds an explicit confidence-and-similarity based order from prior predictions, then refines queries with an RNN along this order to model competition inside the decoder. This ordered recurrent refinement reduces redundant predictions, stabilizes optimization, and improves final detection accuracy. Experiments on multiple DETR-style detectors show consistent gains with comparable efficiency. Code will be released on GitHub.

Multi-agent LLM systems built on frontier backbones often fail to improve reliably as the swarm grows: under strong backbone models, performance frequently plateaus or declines beyond a small number of agents. In a budget-matched controlled intervention on SWE-bench Verified ($N{=}8$, Claude Opus 4.5 backbone, 64-call budget per instance), varying only the communication policy moves resolve rate from 80.4\% under random broadcast coordination (below the 80.9\% single-shot baseline) to 85.5\% under a learned orchestrator, a 5.1-point spread under matched compute. We introduce OrchestraRL, an orchestration layer trained over frozen LLMs. A centralized Macro-Orchestrator reads the swarm's context entropy $H(t)$, computed from the spread of agent output embeddings, and outputs communication mode (via learned entropy thresholds) and topology; per-agent Micro-Coordinators decide whom each agent addresses, when to speak, and at what granularity. Both are trained with REINFORCE and a learned value baseline. Across four benchmarks (SWE-bench Verified, GPQA Diamond, LiveCodeBench v6, SimpleQA), OrchestraRL is comparable to the strongest baseline at $N{=}2$ and achieves the best resolve rate at $N \in \{4, 8\}$ on every benchmark, with the gap widening as $N$ grows. The learned policy recovers interpretable task-dependent communication patterns, including clustered exploration, star-based synthesis, and pipeline reasoning, without these structures being hand-coded.

Pretrained vision transformers provide strong semantic representations, but their token embeddings are not designed to transform predictably under image rotations. This limits their direct adaptation to domains such as satellite imagery, histopathology, microscopy, and other scientific imaging settings, where in-plane orientation is often arbitrary and no canonical pose exists. Existing equivariant adaptation methods for pretrained models typically rely on frame averaging or canonicalization, which can be effective on synthetically rotated natural images but do not teach the adapted model an internal rotation-aware representation. We introduce OrbitLoRA, a lightweight orbit-aware adapter for frozen visual foundation models. OrbitLoRA combines standard LoRA with a typed residual branch that maintains scalar and vector fields over ViT patch tokens. The typed branch is initialized from local steerable image features and harmonic coordinate fields, updated at selected transformer depths, and injected into the frozen token stream through zero-initialized residual gates. During finetuning, paired rotated views supervise the typed state through a transport-and-rotation consistency loss, while the semantic output is trained to remain invariant for the downstream task. This design preserves the rich features of pretrained ViTs while adding a learnable rotation-symmetry prior. Across classification and segmentation tasks, and across DINOv2-, CLIP-, SAM-, and LLaVA-style backbones, OrbitLoRA matches strong LoRA-with-augmentation baselines in high-data regimes and substantially improves adaptation in low-data regimes, where learning the symmetry prior is most valuable.

We study online alignment of large language models under misspecified preference feedback, where the observed preference oracle deviates from an ideal but unknown ground-truth oracle. The online LLM alignment problem is a bi-level reinforcement problem due to the coupling between data collection and policy updates. Recently, the problem has been reduced to a tractable single-level objective in the SAIL (Self-Improving Efficient Online Alignment) framework. In this paper, we introduce a pointwise oracle uncertainty set in this problem and formulate an oracle-robust online alignment objective as a worst-case optimization problem. For log-linear policies, we show that this robust objective admits an exact closed-form decomposition into the original loss function plus an explicit sensitivity penalty. We develop projected stochastic composite updates for the resulting weakly convex objective and prove $\widetilde{O}(\varepsilon^{-2})$ oracle complexity for reaching approximate stationarity.

Regulatory approval of products in high-stakes domains such as drug development requires statistical evidence of safety and efficacy through large-scale randomized controlled trials. However, the high financial cost of these trials may deter developers who lack absolute certainty in their product's efficacy, ultimately stifling the development of `moonshot' products that could offer high social utility. To address this inefficiency, in this paper, we introduce a statistical protocol for experimentation where the product developer (the agent) conducts a randomized controlled trial sequentially and the regulator (the principal) partially subsidizes its cost. By modeling the protocol using a belief Markov decision process, we show that the agent's optimal strategy can be found efficiently using dynamic programming. Further, we show that the social utility is a piecewise linear and convex function over the subsidy level the principal selects, and thus the socially optimal subsidy can also be found efficiently using divide-and-conquer. Simulation experiments using publicly available data on antibiotic development and approval demonstrate that our statistical protocol can be used to increase social utility by more than $35$% relative to standard, non-sequential protocols.

Squared Wasserstein distance is a frequently used tool to compare probability distributions. This distance is typically computed between empirical measures of size $n$ from two underlying random samples. Unfortunately, even in lower dimensional Euclidean space problems $\left( d \in \{2,3\} \right)$, Wasserstein distance algorithms with approximate or exact precision guarantees scale poorly in the runtime as a function of $n$ and the desired precision. In response, we consider the computational-statistical runtime, where the goal to estimate from samples the Wasserstein distance up to the $\varepsilon$-additive error which is achievable under a sample; we allow $O(1)$ time per sample. Towards this, we develop a Sample-Sketch-Solve paradigm where we introduce a regular cartesian grid sketch of the samples. We show that (especially under $\alpha$-H\"older smooth distributions) this can compress the data without increasing asymptotic error, and also regularizes the structure which enables faster exact algorithms. Ultimately, we approximate $W_2^2(P,Q)$ within $\varepsilon$ error in $\varepsilon^{-\max(2,\frac{d+1+o(1)}{2})}$ time for Lipschitz $P,Q$ on $(0,1)^d$; nearly an optimal $\Theta(\varepsilon^{-2})$ for $d \leq 3$.

We demonstrate that AI agents can autonomously improve mathematical proofs of quantitative results. As a key step towards this frontier, we achieve strict improvements on Grothendieck's and Korenblum's constants, whose precise values have remained surprisingly elusive despite their recurrence in both the mathematical and physical sciences. By encoding relevant literature as a partially formalized Lean blueprint, we structure the agent's search space to localized proof refinements that automatically propagate to a verified final bound. Concretely, Grothendieck's constant ($\mathcal{K}_G$) is shown to be: $$ \mathcal{K}_G \leq \tfrac{\pi}{2 \log (1 + \sqrt{2})} - 10^{-17}, $$ which is the first explicit numerical improvement since Krivine's landmark 1979 result. For Korenblum's constant ($\mathcal{K}_K$), the lower bound is tightened to: $$ \mathcal{K}_K \geq 0.3554 + 0.001165, $$ a third-digit improvement over Wang's recent 2025 result. In summary, we establish that AI agents can autonomously refine formal proofs of open problems and provide a generalizable framework for future mathematical discovery.

Tool documentation is the critical interface bridging Large Language Model (LLM) agents and external environments. However, real-world documentation is often ambiguous or insufficient, leading to misaligned tool calls and task failures. While existing methods rely on parameter-efficient tuning or reinforcement learning to enhance tool-use capabilities, they incur substantial computational overhead and lack the agility to adapt to frequently updated toolkits. We propose OpTool , an evolutionary framework that treats tool documentation as an evolvable configuration. OpTool consists of three stages: (1) Contrastive Trajectory Generation, which explores diverse tool-use trajectories via contrastive beam search; (2) Trajectory-based Insight Extraction, which extracts tool-use insights through retrospective global credit assignment; and (3) Tool Documentation Evolution, which iteratively refines the documentation to mitigate information insufficiency. Experimental results across multiple domains demonstrate that OpTool significantly improves tool call accuracy and task success rates, providing a robust solution for tool-based LLM agents.

Mode connectivity has been widely studied, yet the role of the optimizer remains underexplored. We revisit it through optimizer-induced implicit regularization, asking how connectivity behaves when restricted to solutions constrained by a given optimizer. For two-layer ReLU networks, we show that solutions from a single optimizer - AdamW, Muon, or others in the Lion-$\mathcal{K}$ family - form a connected set at sufficiently large width, a result not implied by prior work. We then characterize how optimizer-induced regions interact: at large width two different regions can be disjoint or overlap depending on regularization, while in our small-width example AdamW and Muon converge to disconnected zero-loss components separated by a provable loss barrier. Empirically, in GPT-2 pretraining, we observe same-optimizer paths preserve each model’s spectrum while cross-optimizer paths traverse a smooth transition. Our results reveal optimizer-dependent structure beyond classical mode connectivity literature.