NeurIPS 2026 Papers

Knowledge distillation provides shared semantic supervision for local training in federated learning, reducing the tendency of client models to overfit their own data distributions. Existing methods construct a global knowledge repository from client-uploaded local information to establish consensus across clients. However, global knowledge that is fundamentally derived from client data inevitably contains local biases under heterogeneous settings. To address this issue, this study proposes a geometry-regularized cross-silo consensus enhancement method, termed RISE, which improves collapse resistance in federated learning by preserving pre-trained geometric consensus priors during local adaptation. Specifically, RISE introduces two complementary geometry regularizers. The Biased Low-dimensional Subspace Correction module preserves the spectral energy profile of pre-trained representations by constraining the energy distribution of adapted features within the principal representation subspace. Instead of enforcing direction-wise alignment, it discourages excessive energy concentration along a few dominant components, thereby mitigating biased low-dimensional feature collapse. The Manifold Relation Graph Alignment module regularizes the relational manifold geometry by constructing a transferable inter-class topology from pre-trained global class prototypes and aligning client-side class structures with this global relation, thereby mitigating class representation drift and stabilizing global aggregation. Extensive experiments on eight benchmark datasets across two tasks demonstrate that RISE improves cross-client consensus consistency and enhances the generalization ability of the global model, outperforming eight state-of-the-art methods.

Kolmogorov–Arnold Networks (KANs) replace fixed activations in deep architectures with learnable univariate edge functions, making the choice of edge parametrisation central. Existing variants rely on fixed bases such as splines, polynomials, or Fourier features, which impose a function-space geometry before data are observed. We introduce geometry-constrained KANs, a family of edge activations derived from Banach duality maps in which the geometry itself is learned through a scalar exponent $p > 1$ per edge. This exponent controls the qualitative response: sub-Euclidean values produce sharp threshold-like behaviour, $p = 2$ recovers the linear regime, and larger values produce flatter responses near the origin. Beyond expressivity, the exponent also modulates sensitivity: smaller values reduce amplification of perturbations, providing an implicit regularisation effect without introducing an explicit shrinkage hyperparameter. Across 50 Feynman symbolic regression equations, geometry-constrained KANs match strong fixed-basis baselines on clean data while substantially improving robustness under distribution shift. In particular, they outperform cross-validated spline baselines in 85 of 90 data-scarcity settings and degrade significantly less under noise (3.7× versus 21.6× for splines). Learned exponents are interpretable and stable, revealing consistent geometric structure across equation families and input dimensions.

When language models lack relevant knowledge for a given query, they frequently generate plausible responses that can be hallucinations, rather than admitting being agnostic about the answer. Retraining models to reward admitting ignorance can lead to overly conservative behaviors and poor generalization due to scarce evaluation benchmarks. We propose a post-hoc framework, Conformal Abstention (CA), adapted from conformal prediction (CP) to determine whether to abstain from answering a query. CA provides finite-sample guarantees on both the probability of participation (i.e., not abstaining) and the probability that the generated response is correct. Importantly, the abstention decision relies on prediction confidence rather than the non-conformity scores used in CP, which are intractable for open-ended generation. To better align prediction confidence with the model's ignorance, we introduce a calibration strategy using representation geometry within the model to measure knowledge involvement in shaping the response. Experiments demonstrate that we improve selective answering significantly with 75% conditional correctness.

Zeroth-order optimization (ZOO) is a promising approach for memory-efficient fine-tuning of Large Language Models (LLMs). However, scaling ZOO to larger models on a single GPU necessitates aggressive quantization, which often destabilizes training. A primary cause of this instability is that existing methods treat strictly positive quantization scale parameters as Euclidean variables, potentially leading to boundary violations and optimization difficulties. To address this issue, we propose Hyper-Octant Zeroth-order Optimization (HoZO), a geometry-aware framework that formulates fine-tuning as Riemannian optimization on the positive orthant manifold. By constructing a novel geodesically complete metric and deriving the closed-form exponential map, HoZO inherently enforces positivity constraints without bias-inducing projections. On the theoretical side, we prove that HoZO achieves the optimal oracle complexity for zeroth-order methods. On the empirical side, HoZO outperforms baselines across six downstream tasks and three model sizes. Notably, it achieves a $34\times$ memory reduction compared to first-order fine-tuning, enabling the stable fine-tuning of a 4-bit quantized Llama-2-70b model on a single 48GB GPU.

Similarity measures such as canonical correlation analysis (CCA), representational similarity analysis (RSA) and center kernel alignment (CKA), are widely used to compare the representational geometries used by different neural networks to solve the same task. Yet, because existing methods compare the *extrinsic geometry* of neural representations in state space, rather than their *intrinsic geometry*, they may fail to capture subtle yet crucial distinctions between fundamentally different computations performed by biological and artificial neural networks. Here, we introduce metric similarity analysis (MSA), a novel method which leverages tools from Riemannian geometry to compare the intrinsic geometry of neural manifolds. Within our mathematical framework, we derive several properties of MSA, including coordinate, scale and rotation invariances. We show that MSA can be used to i) disentangle neural computations of deep networks in rich vs. lazy learning regimes, ii) compare the computation-through-dynamics of recurrent neural networks and state-space models during neuroscience working memory tasks, and iii) investigate the statistical manifolds of diffusion models. Overall, we introduce a mathematically grounded and broadly applicable framework to understand the mechanisms behind neural computations by contrasting their Riemannian geometries, with broad applicability to both neuroscience and machine learning.

Precise localization of recording electrodes is fundamental to systems neuroscience; yet, current methods depend on labor-intensive post-hoc histology, which is incompatible with real-time feedback or chronic implants. We test the hypothesis that the electrical signals being recorded themselves carry sufficient anatomical information to localize their recording site in 3D atlas coordinates, without the need for histology. To probe this, we introduce a self-supervised framework that learns geometry-aware representations of single-channel local field potentials (LFP), leveraging probe channel geometry as a self-supervisory signal. We compare speech-based SSL objectives (e.g., masked prediction) against geometry-aware objectives that constrain the embedding space to mirror physical channel layout along the probe. These representations outperform masked-prediction baselines, and the gains compound at larger data scales. We test downstream analysis with 3D coordinate regression and brain region classification, benchmarking three SSL backbones (Wav2Vec 2.0, Whisper, and Data2Vec) against supervised baselines (AnyNet, ViT, and classical spectral classifiers) across datasets, laboratories, species, and probe technologies. Our method achieves state-of-the-art results in brain region classification and 3D coordinate regression from raw single-channel LFP, outperforming all baselines, affirming our hypothesis.

The recent breakthrough Score-Repellent Monte Carlo (SRMC) method improves MCMC sampling by tilting the target density $\pi$ in $\mathbb{R}^d$ with the factor $\exp(-\alpha\,\theta_n^\top s(x))$, where the score function $s(x) = \nabla_x \log \pi(x)$, and $\theta_n$ is the score average from past samples. However, a single scalar repellence strength $\alpha\ge0$ cannot handle both the steep and the flat directions in anisotropic energy landscapes at once. We introduce Geometry-Aware SRMC (GA-SRMC), which generalizes the Euclidean alignment to $\theta^\top M s(x)$ for any symmetric positive-definite matrix $M$. A stochastic-approximation central limit theorem under a trace constraint on $M$ identifies the trace-normalized inverse Fisher $M_{\rm opt}\propto S^{-1}$ as the unique minimizer of the worst-direction covariance bound, elevating the Fisher information $S=\text{Cov}_\pi(s,s)$ from a proposal-side preconditioner (e.g., FisherMALA, natural gradient) to the optimal repulsion-side matrix, with a single online estimate $\widehat S_n$ serving both roles at no additional cost. On an $8$-target simulation spanning the condition number $\kappa\in[1.3,100]$ and dimension $d\in[2,785]$, GA-SRMC dominates SRMC by $4.4\times$ on anisotropic Gaussian ($\kappa=100$) and $6.8\times$ on MNIST under the MALA baseline, by $18\text{ - }32\\%$ on the well-conditioned Bayesian logistic-regression posteriors, and by up to $+48\\%$ under FisherMALA; the advantage scales monotonically with the condition number $\kappa$.

Multi-view 3D reconstruction has achieved remarkable progress with the advent of feed-forward 3D reconstruction models. However, these models are typically trained and evaluated under ideal, degradation-free imaging conditions, whereas real-world observations often contain various degradations that differ significantly from such settings. Improving robustness for multi-view 3D reconstruction under degraded conditions therefore remains an important challenge. We present Geometry-Aware Representation Denoising (GARD), a novel framework that performs diffusion-based multi-view restoration directly in the feature space of a feed-forward reconstructor. This design exploits the geometry-aware feature representations of the reconstructor to effectively recover accurate scene geometry. Furthermore, by employing a decoder, the refined representations can also be used to restore high-quality RGB images, thereby enabling the simultaneous recovery of 3D scene geometry and high-quality imagery. Comprehensive experiments on the DA3 benchmark demonstrate the effectiveness of the proposed GARD framework. Our code and weights will be publicly released for full reproducibility.

Neural operators provide fast surrogates for PDEs but their deterministic predictions limit their use in tasks requiring uncertainty quantification (UQ), especially under geometric variability. Existing approaches primarily model uncertainty in network parameters, largely overlooking the geometry-aware representations learned by the operator itself. We propose REEF-GP (Residual on Embedded Features Gaussian Process), a post-hoc UQ framework that fits a GP to the residuals of a frozen neural operator whose internal embeddings define the kernel feature space. Rather than learning a separate feature map, REEF-GP adapts the operator’s intrinsic coordinate-feature representations to construct geometry-aware uncertainties. To ensure stability and scalability on unstructured domains, REEF-GP incorporates spectral-normalized projections, heteroscedastic geometry-aware noise, and efficient subset-based training that avoids restrictive low-rank approximations. Across five PDE benchmarks with varying geometries, REEF-GP preserves predictive accuracy while achieving calibrated uncertainty estimates competitive with deep ensembles but at a fraction of their cost. Our approach remains robust under geometric distribution shift, with uncertainty concentrating in physically meaningful regions (e.g., shock fronts). Our results demonstrate that accurate and scalable post-hoc UQ for neural operators can be achieved directly in their learned feature space, offering a practical alternative to parameter-centric approaches.

Solving large scale Optimal Transport (OT) in machine learning typically relies on sampling measures to obtain a tractable discrete problem. While the discrete solver's accuracy is controllable, the rate of convergence of the discretization error is governed by the intrinsic dimension of our data. Therefore, the true bottleneck is the knowledge and control of the sampling error. In this work, we tackle this issue by introducing novel estimators for both sampling error and intrinsic dimension. The key finding is a simple, tuning-free estimator of $\text{OT}_c(\rho, \hat\rho)$ that utilizes the semi-dual OT functional and, remarkably, requires no OT solver. Furthermore, we derive a fast intrinsic dimension estimator from the multi-scale decay of our sampling error estimator. This framework unlocks significant computational and statistical advantages in practice, enabling us to (i) quantify the convergence rate of the discretization error, (ii) calibrate the entropic regularization of Sinkhorn divergences to the data's intrinsic geometry, and (iii) introduce a novel, intrinsic-dimension-based Richardson extrapolation estimator that strongly debiases Wasserstein distance estimation. Numerical experiments demonstrate that our geometry-aware pipeline effectively mitigates the discretization error bottleneck while maintaining computational efficiency.

The explosive demand for interactive Large Language Model serving has highlighted the management of the Key-Value cache's dynamic memory footprint as a critical area for performance optimization in inference engines. Modern inference systems overwhelmingly rely on time-centric scheduling heuristics, such as Shortest Job First. However, their theoretical optimality is rooted in traditional schedule modeling, failing to capture the highly dynamic, 2D spatio-temporal geometric growth specific to LLM inference mechanisms. To resolve this, we propose the geometry-aware online scheduling by introducing the Smallest Volume First (SVF) algorithm and its highly efficient variant, 1-bit SVF. Theoretically, we provide a rigorous mathematical foundation for our approach. Utilizing a novel proof methodology, we tighten the worst-case competitive ratio ($\text{CR} \le 48 \rightarrow \text{CR} \le 5$) for SVF with known output lengths. Building upon this core breakthrough, we complete a comprehensive theoretical taxonomy analyzing our algorithms across different traffic scenarios and information availability. Practically, we seamlessly integrate our approach as a plug-and-play layer in vLLM. Extensive evaluations on Llama-3.1 models demonstrate comprehensive performance gains: SVF delivers strong reductions in both average and tail latency, while 1-bit SVF, with merely a single bit information, achieves competitive throughput and latency. This work establishes a theoretically sound and empirically proven approach for resolving memory-constrained scheduling in modern LLM deployments. To facilitate future research, our code is available at \url{https://anonymous.4open.science/r/LLM-inference-scheduling-7970}.

Model merging has emerged as a practical paradigm for integrating multiple independently trained models into a single model without joint retraining. Prior work shows that parameter fusion techniques, such as parameter decomposition, coefficient optimization, and subspace learning, can achieve strong performance while avoiding the cost of joint training. However, these approaches largely treat merging as a parameter aggregation problem, implicitly assuming compatibility in the underlying directional structures. We argue that this assumption rarely holds. In practice, independently trained models often exhibit misaligned dominant directions in both parameter and feature spaces. Naïve merging can therefore disrupt structural coherence, leading to significant performance degradation, especially in heterogeneous settings such as healthcare, where models trained on tasks like medical QA or clinical text understanding develop distinct representational geometries. Moreover, coefficient-based methods rely on well-aligned feature directions, an assumption contradicted by insights from Neural Collapse, which suggest that class features follow structured yet model-specific directional patterns. In this work, we identify directional alignment as the key principle for effective model merging. We propose Merging with Directional Alignment (\method{}), a unified geometric framework that explicitly enforces consistency in directional structures across both parameter and feature spaces. Our analysis shows that directional alignment preserves structural coherence during merging, and extensive experiments across multiple benchmarks, including healthcare QA, demonstrate consistent improvements across model scales and task settings.

Mechanistic interpretability aims to explain a model’s behavior by identifying causally responsible internal structures. Dictionary-based explainers such as sparse autoencoders and transcoders are a primary tool, but their faithfulness under out-of-distribution (OOD) shift has received little systematic attention. We show that distribution shift rotates the subspace that the model actively uses, misaligning the explainer’s dictionary trained on in-distribution (ID) activations. We formalize this misalignment as the faithfulness gap, a geometric distance between the ID dictionary and the OOD-active subspace, and show that it controls OOD faithfulness degradation. To reduce this gap, we propose the Geometry-Adaptive Explainer (GAE), which realigns the explainer's dictionary with the OOD-active subspace while preserving the original feature structure. This requires only unlabeled OOD activations and no gradient updates. We prove that GAE improves over the unadapted ID explainer, with excess loss bounded quadratically by the second-moment shift. Empirically, GAE even matches or surpasses all training-based baselines in causal faithfulness across multiple models and OOD settings.

Out-of-distribution (OOD) detection identifies test samples that fall outside a model's training distribution, a capability critical for safe deployment in high-stakes applications. Standard OOD detectors are trained on a specific in-distribution (ID) dataset and detect deviations from that single domain. In contrast, we study few-shot cross-domain OOD detection: given a \emph{single} pre-trained model, can we perform OOD detection on \emph{arbitrary} new ID-OOD task pairs using only a handful of ID samples at inference time, with no additional training? We propose \textbf{UFCOD}, a unified framework that achieves this goal through information-geometric analysis of diffusion trajectories. Our key insight is that diffusion noise predictions are score functions (gradients of log-density), and we extract two energy features: \emph{Path Energy} (integrated score magnitude) and \emph{Dynamics Energy} (score smoothness), that form a discrete Sobolev norm capturing how samples interact with the learned diffusion process. The central contribution is a \textbf{train-once, deploy-anywhere} paradigm: a diffusion model trained on a single dataset serves as a universal feature extractor for OOD detection across semantically unrelated domains. At deployment, each new task requires only $\sim$100 unlabeled ID samples for inference: no retraining, no fine-tuning, no task-specific adaptation. Using 100 ID samples per task, UFCOD achieves 93.7\% average AUROC across 12 cross-domain benchmarks, competitive with methods trained on 50k--163k samples, demonstrating $\sim$500$\times$ improvement in sample efficiency. See our code in \url{https://anonymous.4open.science/r/UFCOD-22E8/README.md}.

Multimodal Large Language Models (MLLMs) have shown increasing potential for 3D scene understanding and spatial reasoning from videos. However, even with strong visual geometry priors, the hidden semantics of existing 3D-enhanced MLLMs can drift across views, leading to degraded grounding, captioning, and spatial reasoning performance. The underlying reason can be their geometry injection as additional tokens or additive features, which enriches visual content but fails to define how view-dependent semantic states should transform. To address this issue, we propose Lie Routing Vision-Language Model (LieVLM), a geometry-conditioned Lie routing framework that treats geometry as an operator over hidden semantics. LieVLM first elicits local Lie states from 3D geometry tokens by predicting latent rotation states and routing coefficients. It then converts these rotation states into triplet-wise Lie group actions to route projected 2D visual semantics. Finally, a geometry-aware routed fusion module combines the original semantic path, the direct 3D context path, and the Lie-routed transformation path for language generation. This design preserves pretrained visual-language semantics while allowing geometry to actively correct view-induced semantic drift. Extensive experiments on 3D scene understanding and spatial reasoning benchmarks demonstrate that LieVLM improves view robustness, especially in severe-pose-change regimes, and consistently outperforms additive geometry fusion methods. Our results suggest a shift in 3D MLLM design, i.e., geometry should not merely be injected as context, but should operate on semantic representations.

Neural surrogate models for Partial Differential Equations (PDEs) on unstructured 3D geometries are often limited by poor generalization and the high cost of generating large-scale training datasets. Consequently, pre-training has emerged as a critical alternative to enhance the robustness and scalability of these models. In this work, we introduce a pre-training framework tailored to both steady-state and transient regimes. For steady-state problems, we propose a geometry-driven strategy that leverages intrinsic shape descriptors to learn representations of complex 3D domains. For transient problems, we introduce a physics-driven approach based on online generation of synthetic PDE data, enabling scalable pre-training without reliance on expensive datasets. This framework learns transferable features for challenging downstream tasks. Across multiple experiments, our approach achieves up to 3$\times$ faster convergence, 2$\times$ greater data efficiency, and up to 25\% higher accuracy during fine-tuning, particularly under realistic low-data regimes. This methodology provides a practical pathway toward data-efficient neural emulators for large-scale simulations.

Packed high-token 3D fields are not just long-sequence modeling: once a volume is packed into tokens, sequence order is no longer the geometry of the sample. This creates a general failure mode for attention over packed geometric data: sequence-index positional mechanisms can learn the memory layout rather than the object. We introduce Pack3D, a geometry-aware local--global attention operator that separates memory layout from attention geometry. Pack3D keeps local neighborhood attention exact, compresses only distant context into contiguous 3D block summaries, and indexes attention with axial RoPE on true 3D token-grid coordinates rather than packed sequence positions. With RoPE applied before block compression, each query interacts with a geometry-aware mixture of block tokens, so compressed global context remains spatially interpretable after packing. A synthetic repacking benchmark makes the failure mode explicit: learned positional embeddings and packed-sequence RoPE fail under layout changes, whereas axial grid-coordinate RoPE remains invariant. On Allen WTC--11 hiPSC microscopy, our main claim-grade domain, Pack3D improves matched learned-position local--global attention by 0.0978 best macro Dice and 0.1248 final macro Dice at 262k tokens and depth 8. In an Allen-15 joint comparison, Pack3D also improves over a parameter-matched TokenGrid 3D U-Net on both best and final macro Dice. Inference-time ablations show that the trained model genuinely uses its compressed summaries, and the fused forward/backward implementation is $6.49\times$ lower latency than dense SDPA at 65k tokens.

How do language models acquire such a vast array of concepts and abilities from next-token prediction alone? We introduce an interpretability framework that exposes how this single loss relates to different concepts at different points during pretraining. For each layer and checkpoint, we measure the \emph{optimization engagement} with a concept, i.e., alignment between a probe-defined concept subspace and the average gradient outer product (AGOP) of the next-token loss. The concept subspace captures where the concept currently lives in the model's representation; AGOP captures the directions in which the loss exerts the strongest pressure. Across 12 lexical, syntactic, and semantic tasks and 19 Pythia-1B and -410M checkpoints, we find that optimization engagement with a concept is temporally aligned with behavioral changes in the model related to the concept. Furthermore, optimization engagement tends to peak before we observe the largest behavioral changes in the model, suggesting that sudden changes in model performance ("grokking") may be preceded by behaviorally unobservable geometric changes. However, this "geometric precedence" pattern is not universally observed; we also see cases where optimization engagement does not peak once, but rather persists at a high value, or peaks multiple times. These results indicate that the loss contributes to concept learning in a myriad of ways that simple probe performance alone cannot reveal.

Flow matching learns a velocity field whose ordinary differential equation transports a simple reference distribution to a data distribution. Existing theory typically controls this velocity through ambient, isotropic Lipschitz regularity, which becomes extremely pessimistic when data concentrate near a low-dimensional manifold. In this regime, standard error bounds can blow up drastically near the terminal time and fail to explain why flow matching should remain stable. We study flow matching for distributions obtained by smoothing a density on a compact manifold with small Gaussian noise. Our first result is a tube-local Lipschitz regularity theorem showing that, on the relevant neighborhood of the manifold, the population flow-matching velocity has a Jacobian bound that grows only at the inverse effective noise scale. This improves substantially over the stronger singular behavior suggested by existing ambient analyses. Our second result gives a directional description of the velocity Jacobian: tangential directions exhibit mild growth, normal directions are strongly contractive, and tangential-normal coupling is lower order. This reveals that the apparent singularity is highly structured rather than uniformly harmful. These results provide geometric foundations for sharper stability guarantees and suggest a path toward intrinsic-dimension generalization bounds for flow matching on manifold-concentrated data.