NeurIPS 2026 Papers

Robot manipulation under partial observability requires policies to act from incomplete visual evidence while still reasoning about object structure, spatial layout, and physical relations. Existing imitation learning methods mainly encode currently visible observations, which makes them brittle when task-critical parts, contact regions, or support relations are occluded. We introduce Scene Map, a structured scene representation that models object shape, part structure, spatial layout, and physical constraints as a graph of manipulation-relevant primitives and relations. Scene Map turns occluded but task-relevant structure into explicit prediction targets and provides structural priors for inferring unseen components from partial observations. Building on Scene Map, we propose MapPolicy, an imitation learning framework that fuses structured map features with visual and robot state features for action prediction. MapPolicy uses structure-modulated graph attention to propagate global structural and spatial information through primitive features, and a physical constraint loss to regularize learned features with geometric and kinematic consistency. Across 49 manipulation tasks in three simulation benchmarks and five real world tasks, MapPolicy consistently improves over strong imitation learning baselines, with especially clear gains under severe occlusion. Our code will be released publicly to facilitate reproducibility and further research.

Feature caching is a promising way to accelerate Diffusion Transformers (DiTs) without modifying the model backbone, but existing methods largely rely on local heuristics or input-agnostic rules when deciding which intermediate features to cache and reuse. We present Map-Guided Caching, a feature caching framework that predicts a Temporal Variation Map (TVMap) for each denoising trajectory and uses the map to derive a global caching policy. The TVMap summarizes block-wise temporal variation across timesteps, while a lightweight Map Predictor estimates this variation from the prompt and initial noise before denoising starts. Given the predicted map, we formulate policy generation as a global path planning problem and combine a dynamic-programming-based orientation planner with a quantity planner that targets a desired speedup ratio. Across image generation, image editing, and video generation, the proposed method achieves up to $3.3\times$ speedup on FLUX.1 [dev] and HunyuanVideo, and $2.7\times$ on FLUX.1 Kontext [dev], while preserving competitive output quality. The advantage is especially pronounced on complex prompts, where local caching policies degrade more sharply. These results suggest that input-adaptive, globally planned caching is an effective alternative to purely local feature reuse policies.

Modeling fine-grained human manipulation of object functional parts is essential for achieving human-like dexterous manipulation in virtual reality and robotics. Despite recent advances, data-driven generation methods remain largely constrained by patterns learned from limited training data, rendering them prone to hallucinating task-inconsistent manipulations when encountering samples beyond the training distribution. To address this challenge, we propose ManipulationRAG, a framework that augments an internal generative model with externally retrieved, task-relevant manipulation knowledge. Specifically, we design a hand-centric manipulation taxonomy and construct a structured knowledge base that organizes task descriptions, objects, manipulation types, and corresponding hand pose sequences, supported by a retrieval mechanism for accurate knowledge retrieval. The retrieved knowledge serves as an external manipulation-specific prior, guiding the diffusion model toward task-consistent and physically plausible manipulation synthesis. Extensive experiments demonstrate that our method outperforms existing methods and generalizes robustly to object manipulation beyond the original diffusion model's training domain.

We establish universal approximation for neural networks whose weight matrices take values in matrix manifolds. The nonlinear matrix manifolds considered in this work arise as orbits of classical Lie group actions embedded in the ambient space $\mathbb{R}^{d \times k}$. Stiefel manifolds provide a prominent example with empirical precedent in prior neural-network models. We present an expanded catalog of such manifolds, drawing from well-studied constructions in physics and differential geometry. These manifolds are particularly appealing for intrinsic Riemannian optimization: the transitivity of the underlying group action yields explicit tangent-space descriptions and retractions. Moreover, the group-action structure enables parameter-efficient factorizations, reducing the number of trainable variables. This leads to a fundamental expressivity question: when weight matrices take values in nonlinear matrix manifolds, which architectures retain universal approximation power? We prove that manifold-weighted networks with residual connections and bounded scalar multipliers are $L^p(K)$-universal. To establish universality across our manifold catalog, we introduce two independent sets of sufficient conditions: Register-Isolated Primitive Realizability (RIPR) and Cross-Axis Fold-and-Cut (CAFC). These conditions correspond to two distinct approximation constructions. By verifying RIPR or CAFC case by case for the matrix manifolds in our catalog, we establish universality for the corresponding residual manifold-weighted networks with bounded scalar multipliers.

Dense stereo disparity estimation provides the geometric basis for 3D perception in computer-assisted endoscopy. Endoscopic stereo matching is difficult because reliable tissue correspondences often appear next to weak-texture, specular, occluded, or deformed regions within the same frame. Existing stereo and stereo-monocular methods improve cost aggregation or introduce monocular priors, but they often treat monocular guidance as a global feature source or an independently predicted candidate. Without evaluating the local support of the current stereo match during refinement, the update can reinforce corrupted correspondences in ambiguous regions or allow monocular priors to override reliable metric matches. To tackle these problems, we introduce Manifold-Stereo, a reliability-gated latent-flow framework for endoscopic stereo disparity estimation. The method refines a monocular disparity candidate through a conditional latent residual flow on an image-aligned correction space induced by stereo cost-volume geometry, then calibrates the corrected candidate to the current stereo disparity scale. A geometric consistency gate computed from left-right feature alignment regulates the candidate contribution during recurrent refinement, preserving stereo estimates in well-matched regions while supplying monocular context where local matching is ambiguous. Experiments on in-domain SCARED and zero-shot SERV-CT and Hamlyn show strong average disparity accuracy and competitive thresholded outlier rates against recent stereo and endoscopic baselines.

Anti-customization techniques aim to protect visual content from unauthorized diffusion model personalization by adding imperceptible perturbations to training images. However, such protection can be significantly weakened by diffusion-based purification, which may remove adversarial perturbations before customization. We study representative anti-customization methods under diffusion purification and observe that existing methods exhibit similar off-manifold perturbation structures that are largely suppressed by purification. We then provide a principled geometric analysis that characterizes purification as a local contraction toward the data manifold, primarily suppressing perturbation components along normal directions. Motivated by this insight, we propose Manifold-Aligned Customization protection **MAC**, a practical purification-aware framework that optimizes downstream anti-customization objectives through a differentiable purifier and uses a structured perturbation parameterization to bias updates toward purification-retentive, manifold-aligned directions. Across experiments on CelebA-HQ, VGGFace2-HQ, and WikiArt under diffusion-based purification, MAC substantially improves protection effectiveness across most prompts and domains. Finally, frequency analysis and ImageNet and MS-COCO retention studies show that spectral signatures cannot predict purification outcomes, providing additional evidence for the proposed manifold-alignment perspective beyond a single data domain.

We present a new paradigm for creating random features to approximate bi-variate functions (in particular, kernels) defined on general manifolds. This new mechanism of $\textit{Manifold Random Features}$ (MRFs) leverages discretization of the manifold and the recently introduced technique of $\textit{Graph Random Features}$ (GRFs) to learn continuous fields on manifolds. Those fields are used to find continuous approximation mechanisms that otherwise, in general scenarios, cannot be derived analytically. MRFs provide positive and bounded features, a key property for accurate, low-variance approximation. We show deep asymptotic connection between GRFs, defined on discrete graph objects, and continuous random features used for regular kernels. As a by-product of our method, we re-discover recently introduced mechanism of Gaussian kernel approximation applied in particular to improve linear-attention Transformers, considering simple random walks on graphs and by-passing original complex mathematical computations. We complement our algorithm with a rigorous theoretical analysis and verify in thorough experimental studies.

Hyperspectral Image (HSI) reconstruction has made substantial progress with the deep unfolding framework by decomposing the problem into a projection module and a denoiser module. Nevertheless, existing methods still exhibit limitations of insufficient matching with HSI data. The issues lie in two aspects: 1) the data-consistency projection module applies a fixed gradient descent step while ignoring the spectral content adaptivity of HSI; 2)the denoiser module either suffers from high computational complexity or is limited to the local receptive field, failing to efficiently capture the global spatial-spectral dependencies. In this paper, we propose a deep unfolding framework centered on a two-stage manifold learning strategy to obtain a degradation-free structural prior, which is then strategically injected into both the projection and denoiser modules. In the data-consistency projection, a prior-guided spectral weight mechanism is introduced to adaptively modulate the weight of spectral band. In the denoiser module, a manifold-guided Mamba is proposed to achieve the structure-aware extraction of spatial-spectral features through prior modulation. Quantitative and visual comparisons on synthetic and real-world datasets show that the proposed method significantly outperforms existing state-of-the-art approaches, while maintaining the structural interpretability of the unfolding framework. To support reproducibility, the code will be publicly released upon acceptance.

Estimating semantic correspondences between different object instances of similar categories in real-world images is a fundamental challenge in computer vision. While the recent progress in foundation models has significantly advanced solutions for such difficult image correspondence problems, respective methods are often insufficiently regularised. For example, common nearest-neighbour matching with foundation model features often results in noisy correspondences, which is particularly prominent when solving for dense correspondences. In this work, we tackle image matching via functional maps, which have been popularised for 3D shape matching due to their powerful spectral formalism utilising an efficient low-dimensional linear representation. However, functional maps require that domains (in our cases the rectangular image grids) are equipped with an informative and non-trivial geometry. To define a meaningful manifold structure for images, we use Galerkin's method from finite elements to derive a discrete Laplace-Beltrami operator (LBO) based on a manifold embedding of deep image features. Further, we leverage a neural adjoint map that allows to represent non-linear mappings. We demonstrate that our method sets the new state of the art among zero-shot image correspondence methods on multiple dense correspondence benchmarks of real-world images.

Preference optimization is a standard alignment method for generative models, yet extending it to continuous-time dynamics remains non-trivial. In flow matching, reward-driven updates modify the transport trajectories but lack inherent constraints to pretrained data manifold, often pushing generated (terminal) samples off the pretrained support. We formalize this failure mode as ``manifold drift``, a phenomenon where preference-aligned trajectories diverge from the pretrained support. Theoretically, we prove that while optimal flow matching exactly preserves the terminal manifold, existing methods like FlowDPO fail to do so whenever reward-driven updates contain components normal to the manifold surface. As a remedy, we propose **ThermoDPO**, a temperature-controlled objective that anchors pairwise preference optimization on preferred samples. Across temperature regimes, this objective connects rejection sampling fine-tuning and FlowDPO while upper-bounding a reconstruction-based surrogate for manifold drift. To counteract diminished signals at low temperatures, we further introduce a weighted variant, **ThermoDPO-weighted**, which we evaluate across all experiments for its superior trade-off between preference alignment and manifold preservation. Empirically, **ThermoDPO-weighted** demonstrates a better trade-off on synthetic manifolds than FlowDPO variants. On real-world image benchmarks, it improves OCR-oriented alignment over the pretrained model and FlowDPO variants while remaining competitive with strong baselines in held-out model-based and human evaluations without visible sample degradation.

This paper addresses the problem of sequential decision-making under learning budget constraints. Such settings naturally arise in applications like managing a portfolio of bandit or reinforcement learning (RL) algorithms. We propose a novel UCB-type algorithm, M-LCB, designed to manage a pool of $K$ self-learning experts in a stochastic environment while accounting for a limited per-round learning budget $M$. At each round, M-LCB selects one expert to make a decision and at most $M \le K$ experts to learn. For selection, M-LCB uses confidence bounds constructed from limited prior knowledge about the experts (i.e., mild assumptions) and their observed training losses. We derive anytime regret bounds for M-LCB that scale with the individual regrets of the experts. In particular, if each expert has regret $\tilde O(T^\alpha)$ by round $T$, then M-LCB guarantees an overall regret of $\tilde O\left(\sqrt{KT/M} + (K/M)^{1-\alpha}T^\alpha\right)$ relative to the best expert in hindsight. Finally, we demonstrate the applicability of M-LCB using self-learning experts instantiated as (i) parametric models and (ii) bandit algorithms.

Vision-language models (VLMs) can achieve strong image-level visual question-answer (QA) performance while relying on shortcuts rather than grounded visual understanding. This is especially concerning in medical imaging, where clinically meaningful interpretation often requires spatial reasoning, localization, and anatomical understanding. Yet in mammography, no benchmark directly assesses whether VLMs can localize findings or reason spatially. To address this, we introduce MammoGPS, a benchmark for evaluating visual grounding, perception, and spatial reasoning in mammography. MammoGPS comprises 56,638 QA pairs across 13 clinical tasks for anatomical and finding grounding, finding-type, quadrant prediction, and depth determination. MammoGPS also includes 94,119 QA pairs across 20 controlled diagnostic tasks designed to probe model failures and spatial biases: we design controlled difficulty ladders to isolate the effects of semantic cues, visual cues, clinical terminology, and background context, as well as spatial-bias tasks to test whether model performance depends on finding location. We evaluate 12 VLMs spanning proprietary, general-purpose, and medical-domain models. Our analysis shows that current systems struggle to localize suspicious findings, make limited use of spatial semantic cues, rely on priors and shortcut heuristics for categorical and grounding tasks, struggle with mammography-specific spatial terminology, and exhibit position-dependent biases. Together, our results underscore the need for more rigorous and fine-grained mammography VLM evaluations.

Spatial interpolation from sparse irregular observations is a recurring task across the environmental and health sciences. Neural processes amortise this inference into a single forward pass at test time. Joint-predictive variants capture correlations across query locations but typically rely on self-attention over context and query points, incurring $O(K^2)$ compute in the combined point count $K$ and capping practical use at a few thousand points. State-space models offer a linear-time alternative, but existing state-space NPs compress the context into a single hidden state before reading queries, reducing the information available to the decoder. We propose MambaFlowNP, a joint-predictive NP with linear compute trained with a conditional flow-matching objective. A multi-directional Mamba-2 scan reads context and query tokens in a single shared sequence, so each query's velocity is read from its own scan position rather than from a compressed context summary. MambaFlowNP is competitive with the strongest attention-based joint-predictive baselines on synthetic Gaussian-process benchmarks, and across regional and continental-scale NOAA temperature interpolation tasks it attains the lowest RMSE on every tier and the lowest CRPS at the California and continental US tiers, while running $170$--$760{\times}$ faster than the strongest joint-predictive baseline. A single forward pass scales to $10^6$ points on one GPU in under a second. This scalability unlocks amortised joint-predictive inference for continental-scale, irregularly sampled sensor networks, spanning climate monitoring to disease vector surveillance.

Mamba, a recently proposed linear-time sequence model, has attracted significant attention for its computational efficiency and strong empirical performance. However, a rigorous theoretical understanding of its underlying mechanisms remains limited. In this work, we provide a theoretical analysis of Mamba's in-context learning (ICL) capability by focusing on tasks defined by low-dimensional nonlinear target functions. Specifically, we study in-context learning of a single-index model $y \approx g_*(\langle \boldsymbol{\beta}, \boldsymbol{x} \rangle)$, which depends on only a single relevant direction $\boldsymbol{\beta}$, referred to as *feature*. We prove that Mamba, pretrained by gradient-based methods, can achieve efficient ICL via *test-time feature learning*, extracting the relevant direction directly from context examples. Consequently, we establish a test-time sample complexity that improves upon linear Transformers---analyzed to behave like kernel methods---and is comparable to nonlinear Transformers, which have been shown to surpass the Correlational Statistical Query (CSQ) lower bound and achieve near information-theoretically optimal rate in previous works. Our analysis reveals the crucial role of the *nonlinear gating* mechanism in Mamba for feature extraction, highlighting it as the fundamental driver behind Mamba’s ability to achieve both computational efficiency and high performance.

Graph Neural Networks (GNNs) have shown remarkable effectiveness across diverse graph-related tasks. However, existing studies reveal that their fairness is highly vulnerable to adversarial manipulation, which can substantially exacerbate inherent biases toward sensitive attributes, such as gender in demographic prediction. Although prior work has shown that fairness poisoning attacks can be launched via malicious node injection, whether this strategy can support the more challenging fairness evasion attacks remains open. Moreover, existing methods have not directly addressed fairness attacks on multi-class datasets with multi-valued sensitive attributes. To bridge this gap, we propose a universal fairness attack framework (UFA) that injects malicious nodes to perform both poisoning and evasion attacks against GNNs while minimally affecting model utility. UFA first identifies nodes most susceptible to being pushed across the decision boundary. It then targets core mechanisms shared by diverse GNN architectures to manipulate node representations, thereby maximizing outcome disparities across sensitive groups. Comprehensive experiments on five real-world datasets demonstrate the effectiveness of UFA. By injecting fewer than 0.2\% malicious nodes, UFA severely degrades the fairness of both mainstream and fairness-aware GNNs, proving effective in both poisoning and evasion settings. Our findings expose the fragility of fairness in GNNs and underscore the urgent need for robust fairness-aware models.

Open-source LLMs (OSMs) are reaching near state-of-the-art performance, prompting prior works to trace the text they generate by embedding text watermarking algorithms directly into their weights. Yet, OSMs are subject to post-training modifications, which has been shown to remove the watermark. Model merging in particular, a prominent method used for combining expert knowledge and preventing catastrophic forgetting, strongly removes such OSMs watermarks. A key question is how to enable OSM watermarks that survive subsequent merging. In this work, we show for the first time how to design an OSM watermark that is durable against model merging. We propose Merge-Adversarial Training, an adversarial training algorithm to distill text watermarks into model weights while being robust to subsequent model merging. Our approach consistently outperforms all baselines (e.g. with SLERP up to +51 percentage points (pp) TPR@1%FPR with +25 pp on average) while preserving downstream capabilities. We also for the first time evaluate OSM watermarks against realistic merge scenarios, representing common use-cases such as combining expert capabilities or preventing catastrophic forgetting, and with 3 prominent merging algorithms. More broadly, our findings suggest that adversarial training is a reliable approach for increasing OSM watermark durability against post-training modifications.

The key--value (KV) cache is a major bottleneck in long-context inference, where memory and computation grow with sequence length. Existing KV eviction methods reduce this cost but typically degrade performance relative to full-cache inference. Our key insight is that full-cache attention is not always optimal: in long contexts, irrelevant tokens can dilute attention away from useful evidence, so selective, learnable eviction can improve generation rather than merely approximate the full cache. We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget. Lightweight retention gates assign utility scores to cached KV entries, and a shared final scoring projection calibrates these scores across all layers and heads. This enables a single global eviction policy in which tokens from different layers, heads, and modalities compete directly for cache capacity. We further provide theoretical analysis showing that preferentially retaining useful tokens reduces attention dilution, and we justify geometric retention as a query-agnostic proxy for future utility. Across long-context language and vision--language benchmarks, our method substantially reduces KV memory while matching or surpassing full-cache inference. These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.

Finding the optimal sample complexity of realizable PAC learning was a major open problem in learning theory, which was settled by the breakthrough result of Hanneke (2016), building on prior work by Simon (2015). Hanneke's algorithm is quite involved and requires training $\text{poly}(n)$ Empirical Risk Minimizers (ERMs) on carefully crafted subsets of the data. Since using a single ERM call is provably suboptimal, recent work by Aden-Ali (2024) proposed the majority-of-three ERMs as a candidate for the *simplest* optimal PAC learner, and established its optimality in the *in-expectation* regime. Their main open question was whether this algorithm is optimal in the classical *high-probability* regime. In this work, building on their ideas, we answer this question affirmatively.

The growing deployment of Large Language Models (LLMs) has raised concerns about their misuse in generating harmful or deceptive content. To address this issue, watermarking methods have been proposed to embed identifiable multi-bit messages into generated text for misuse tracing. However, existing methods often suffer from a fundamental trade-off between text quality and decoding accuracy. In particular, they have to restrict the size of the preferred token set (i.e., green list) during encoding to maintain a detectable watermark signal for decoding, which inevitably degrades generation quality. To improve this trade-off, we propose a novel message encoding paradigm called \textit{majority bit-aware encoding}, which relaxes the watermark signal strength from the green list size. This strategy allows for a strong watermark signal to be preserved in generated texts even when using a large green list. We introduce two instantiations of this paradigm: MajorMark and MajorMark$^{+}$, where the latter is specifically optimized for long messages. Extensive experiments on state-of-the-art LLMs demonstrate that our methods achieve higher decoding accuracy and superior text quality compared to prior baselines. The code of the proposed methods is available \href{https://anonymous.4open.science/r/MajorMark}{here} for review.