NeurIPS 2026 论文列表

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.

Learning the natural parameters $z \in \mathbb{R}^n$ of discrete distributions $\mu_z$ from independent samples constrained to a subset $S \subseteq \\{0,1\\}^n$ is a foundational challenge in high-dimensional statistics. Existing methods for efficiently estimating truncated Boolean product distributions, notably the work of [Fotakis et al' COLT'20, Algorithmica '22], require either strong local connectivity assumptions on $S$ -- a property denoted *fatness* -- or stringent anti-concentration assumptions and necessitate the total mass of the truncation set to be a constant with respect to $n$. Moreover, the results in [Fotakis et al' COLT'20, Algorithmica '22] suffer from sample complexities that scale as $\Omega(2^n)$ if the mass of $S$ is exponentially small in $n$. In this work, we circumvent these limitations by analyzing the geometry of $S$ under the measure $\mu_z$. We refine the existing parameter estimation guarantees under the fatness assumption, improving the prior sample complexity to $\mathcal{O}( \log n / \epsilon^2)$ for $\ell_\infty$-recovery, matching the untruncated minimax rate. We further generalize fatness using the notion of influence utilized in the analysis of Boolean functions and provide sufficient conditions for efficient inference. Notably, unlike previous work, our method does not require sampling at arbitrary parameterizations of the model and instead relies on gradient descent. Lastly, we establish a theoretical lower bound demonstrating the sample complexity exhibits an intrinsic exponential dependence on the width of the model and the minimum distance between elements in the set.

Large language model (LLM) inference serving is undergoing rapid growth and large-scale deployment, motivating us to rethink how the inference process itself can be leveraged to enable efficient task-specific LLM adaptation. In this paper, we propose a forward-only approach for efficient LLM adaptation with forward-only passes in LLM inference serving. We exploit the angle concentration of activations induced by each singular value decomposition (SVD) component to measure its contribution—_dispersed_ or _concentrated_—to the angle concentration of the hidden states. Based on this, our forward-only search then efficiently identifies the weight matrix with the highest dispersion that merits rank reduction. We then selectively remove the higher-order components and retain the lower-order components in SVD. Empirical results across diverse datasets demonstrate the competitive accuracy of our forward-only approach, while theoretical analysis shows lower peak memory and greater speedup than the gradient-based approach, and greater speedup than the exhaustive search approach. The extended experiments further present the robustness and generalization of our forward-only approach to various LLMs with up to 57B parameters. The code is included in the supplementary material.

Structured sparsity extends classical sparse learning by encouraging sparsity at the level of predefined groups rather than individual parameters. This paradigm is widely adopted in modern machine learning and is typically enforced by sparsity-inducing norms. In real pipelines, the optimal group partition is often unknown a priori and does not remain static after initial training. Parameter groupings are often modified as domain knowledge or upstream representation evolves. Standard approaches require retraining the model from scratch under the new target partition, which discards the optimization effort already invested and becomes computationally prohibitive at scale. In this work, we focus on the group-wise $\ell_{p,1}$-norm, a flexible regularization family that encompasses classical group Lasso and other sparse variants. We develop a principled algorithmic framework that leverages the pre-trained optimum to fine-tune model weights toward a new desired group structure without full retraining. Technically, we decompose arbitrary group-wise transitions into a sequence of tractable primitives and derive the closed-form learning dynamics for each sub-action. Extensive experiments on real-world datasets demonstrate that our method achieves significant acceleration while provably recovering the correct optimal solution.

Pre-trained diffusion and flow policies have emerged as powerful backbones for robotics, yet adapting them to novel tasks via reinforcement learning (RL) remains computationally unstable, sample inefficient and lacks theoretical guarantees. Prior approaches that fine-tune the latent noise space bypass costly backpropagation through time but struggle with "noise aliasing," where determining the optimal noise input to the pre-trained generative model requires repeatedly refitting separate Q-value networks. This paper resolves these challenges by introducing a novel theoretical framework demonstrating that imposing a linear structure on the environment induces a Latent Noise Linear Markov Decision Process (LNL-MDP). By treating the pre-trained model as a virtual dataset, our LNL-MDP fine tune framework falls under the FineTuneRL theory with near-optimal guarantee. To translate this idealized theory into a practical methodology, we derive the Bridge Equation which eliminates the noise aliasing problem. Leveraging this insight, we propose Latent Representation Bridging for Diffusion (LARBRID), an algorithm that efficiently learns the representation space from the LNL-MDP and optimal noises to steer the pre-trained model to act optimally in given tasks. We demonstrate that LARBRID is exceptionally sample efficient, yielding up to $2\text{-}3\times$ performance increases on challenging control problems across a diverse suite of robotic locomotion and manipulation tasks.

Collaboratively fine-tuning (FT) large language models (LLMs) over heterogeneous mobile devices fosters immense potential applications of personalized intelligence. However, such a vision faces critical system challenges. Existing federated LLM fine-tuning approaches remain limited by prohibitive on-device resource overhead, straggler-prone synchronous aggregation, and the assumption of a unified pretrained backbone across heterogeneous devices. In this paper, we propose EC-MobiLLM, a novel design for efficient collaborative LLM fine-tuning across heterogeneous mobile devices and heterogeneous backbones. EC-MobiLLM adopts a server-assisted side-tuning paradigm to minimize on-device overhead and pioneers non-blocking asynchronous collaboration to accelerate training. Moreover, EC-MobiLLM introduces adaptive feature alignment to support collaboration across heterogeneous backbones, aligning with practical needs in real-world deployments. Extensive experimental results demonstrate that EC-MobiLLM can maintain robust fine-tuning performance while achieving extremely low on-device memory, with at least 95.2\% reduction in computation overhead, 93.2\% reduction in communication costs and $5.1\times$ faster convergence compared to existing methods, validating its efficacy for practical LLM adaptation over heterogeneous mobile devices.

We consider the problem of maintaining a minimum-cost bipartite matching between two point sets $A$ and $B$ in a metric space $(\mathbb{X},\mathsf{d})$ under insertions and deletions of input points. The Wasserstein-$p$ ($W_p$) cost of a matching $M$ is defined as $(\sum_{(a,b)\in M}\mathsf{d}(a,b)^p)^{1/p}$ and an $\alpha$-approximate matching is one whose total cost is at most $\alpha$ times the cost of a minimum-cost matching. We obtain two main results. If $A$ and $B$ are points in $\mathbb{R}^d$ and the distance between two points is measured under the $\ell_p$-metric, an $O(d\epsilon^{-3/2})$-approximate $W_1$-matching between $A$ and $B$ can be maintained with an amortized $O\big(n^\epsilon d\epsilon^{-3/2}\varphi(n,1/\sqrt{\epsilon})\big)$ update time, for any parameter $\epsilon\in(0,1]$, and where $\varphi(n,1/\sqrt{\epsilon})$ denotes the update time of a $\epsilon^{-1/2}$-approximate nearest neighbor data structure. For the $\ell_2$ norm this results in an update time of $O\big((d+ \epsilon^{-3/2}n^\epsilon)\log n\big)$. The current state-of-the-art for high-dimensional settings only focus on the $\ell_2$-metric and maintain an $O(\log n)$-approximate matching. Our result not only extends to any $\ell_p$ norm, but it also improves the approximation factor for $d=o(\log n)$ under the $\ell_2$-metric. Next, we consider the problem of maintaining a $W_p$-matching when $A$ and $B$ are point sets in an arbitrary finite metric space. We show that for any $k,p\geq 1$ and constant $\epsilon>0$, a $2k(1+\epsilon)$-approximate matching can be maintained with an amortized update time of $O(kn^{1+1/k}\epsilon^{-1}\log \Delta)$, where $\Delta$ is the spread of $A\cup B$. Prior work on finite metric spaces gave an insertion-only data structure for $W_1$-matching. In contrast, we develop a fully dynamic approach that supports both insertions and deletions and work for all $W_p$-matchings. Together, our results extend dynamic Wasserstein matching beyond fixed-dimensional $W_1$ settings: they handle high-dimensional geometric instances, extend prior insertion-only $1$-Wasserstein matching guarantees to the fully dynamic setting, and support $p$-Wasserstein costs for every integer $p\ge 1$ in arbitrary metric spaces.

Knowledge editing has become a key technique for updating knowledge in large language models. However, existing methods neglect the geometric structure of the knowledge representation space, leading to two key limitations: they cannot distinguish which pre-trained knowledge should be preferentially preserved (core vs. non-core) nor which new edits are inherently difficult (hard vs. easy). To address this, we propose \textsc{GEdit}, a \textit{geometry-aware} framework built on a Low-Rank Plus Diagonal (LR+D) factor model that captures the intrinsic low-rank structure of Transformer FFN hidden representations. \textsc{GEdit} features two complementary mechanisms. First, the \textit{low-rank regularization} applies targeted protection on a low-rank core knowledge subspace. Second, the \textit{adaptive regularization} dynamically adjusts the regularization strength based on each edit request's editability: relaxing constraints for hard edits to ensure success, and tightening them for easy edits to better preserve pretrained knowledge. Extensive experiments on \texttt{GPT2-XL}, \texttt{GPT-J}, and \texttt{LLaMA-3-8B} under large-scale editing scenarios—including both mini-batch and fully sequential editing settings with up to 10,000 continuous updates—show that \textsc{GEdit} consistently and significantly outperforms state-of-the-art baselines in both editing success and knowledge preservation.

Test-time model evolution offers a promising way for deployed models to improve from unlabeled test-time experience, yet most existing methods depend on backpropagation (BP), which incurs substantial memory overhead and makes them difficult to deploy on edge devices, quantized models, specialized accelerators, or black-box models. In this work, we study test-time model evolution under a strict two-forward budget, a setting that pushes adaptation toward highly efficient real-world deployment. We reveal three key obstacles in zeroth-order test-time optimization: susceptibility to shortcut solutions, uncontrolled weight drift, and ineffective update direction estimation. To overcome them, we propose EVA-0, a minimal zeroth-order adaptation framework that: 1) keeps the loss scale-invariant to prevent shortcut solutions; 2) devises an anchor-guided optimization strategy to alleviate weight drift; 3) uses sample-wise symmetric two-sided perturbation for update direction estimation and inference. EVA-0 requires no BP and performs both inference and adaptation within only two forward passes per sample. Results on ImageNet-C\&ViT-Base show that EVA-0 outperforms both BP-based DeYO and BP-free FOA, while achieving a 14$\times$ speed-up over FOA.

Reinforcement learning (RL) has become a key training step for improving mathematical reasoning in large language models (LLMs), but it often has high GPU memory usage, which makes it hard to use in settings with limited resources. To reduce these issues, we propose Evolution Strategies with Sharpness-Aware Maximization (ESSAM), a full parameter fine-tuning framework that tightly combines the zeroth-order search in parameter space from Evolution Strategies (ES) with the Sharpness-Aware Maximization (SAM) to improve generalization. We conduct fine-tuning experiments on the mainstream mathematica reasoning task GSM8K. The results show that ESSAM achieves an average accuracy of 78.27\% across all models and its overall performance is comparable to RL methods. It surpasses classic RL algorithm PPO with an accuracy of 77.72\% and is comparable to GRPO with an accuracy of 78.34\%, and even surpassing them on some models. Further generalization experiments show that the models trained with ESSAM exhibit stronger generalization ability. Their average performance achieves the best results on 5 out of 6 datasets, indicating that ESSAM can effectively improve the generalization performance of fine-tuned models. In terms of GPU memory usage, ESSAM reduces the average GPU memory usage by $18\times$ compared to PPO and by $10\times$ compared to GRPO, achieving an extremely low GPU memory usage. In addition, we design an accelerated variant of ESSAM, which achieves nearly a twofold speedup while maintaining the same GPU memory usage as ESSAM, and attains an average accuracy of 78.02\% across all models, outperforming PPO. Code: https://anonymous.4open.science/r/ESSAM-3F4F/

Large language models are remarkably capable, yet how computation propagates through their layers remains poorly understood. A growing line of work treats depth as discrete time and the residual stream as a dynamical system, where each layer's nonlinear update has a local linear description. However, previous analyses have relied on scalar summaries or approximate linearizations, leaving the full spectral geometry of trained LLMs unknown. We perform full Jacobian eigendecomposition across three production--scale LLMs and show that training installs a monotonic spectral gradient through depth---from non-normal, rotation-dominated early layers to near--symmetric late layers---together with a cumulative low-rank bottleneck that funnels perturbations into a small fraction of the residual stream's effective dimensions. Our experiments reveal that this gradient and the dimensional collapse are learned rather than architectural, and is largely dissolved when structured non-normality is removed. We further show that the topological positioning of graph communities predicts whether the Jacobian amplifies or suppresses them, with the sign of the coupling determined by the local operator type, a relationship absent at initialization. These results map a learned spectral geometry in LLMs that links perturbation propagation and compression to the network's functional topology.

Model merging combines fine-tuned models into a unified one without joint training or access to original data. Dynamic merging improves flexibility by selectively activating task-relevant parameters, but existing methods either maintain a full shared model with tiny experts or allocate excessive capacity to experts, leading to suboptimal accuracy–efficiency trade-offs. We propose **DiDi-Merging**, a slim dynamic merging framework that uses **differentiable rank allocation** to balance shared and per-task expert parameters. The method casts parameter budgeting as differentiable rank optimization over low-rank modules, supervised data-free by the original task vectors, and applies a brief refinement stage to recover task fidelity. Across vision, language, and multimodal benchmarks, DiDi-Merging matches prior dynamic baselines at **1.24×** the parameters of a single fine-tuned model and surpasses them at **1.4×**, substantially more compact than methods requiring ≥2× storage.

Large-scale foundation models achieve strong performance across diverse tasks, but their size makes inference costly, largely due to dense matrix multiplications. Prior work reduces this cost by replacing dense weight matrices with efficient structured forms such as low-rank factorizations. However, these methods approximate weights rather than the output activations that determine inference accuracy. Consequently, small weight-space errors can be amplified by input activations, producing large output errors. In this work, we propose DyRA, an input-adaptive method that improves structured matrix multiplication approximation by correcting residual output errors during inference. We show that matrix multiplication can be approximated more effectively by directly optimizing low-rank factors of the output. DyRA builds on this insight by dynamically approximating and correcting the output error introduced by structured weight approximations. This combines efficient structured computation with input-dependent correction, yielding a more faithful approximation of full matrix multiplication under the same computational budget. Across vision, speech, and language models, DyRA consistently improves the accuracy–efficiency trade-off over structured weight approximations alone. Notably, DyRA achieves a $1.5\times$ end-to-end GPU speedup for DINOv3 while reducing accuracy degradation by more than $3\times$ relative to weight-only baselines.

Adapting Multimodal Large Language Models (MLLMs) to specialized and diverse downstream applications necessitates fine-tuning, typically through parameter-efficient methods such as *Low-Rank Adaptation* (LoRA). However, fine-tuning MLLMs severely degrades pre-trained upstream capabilities, a phenomenon known as *catastrophic forgetting*. Existing Methods primarily rely on parameter pruning to strike a balance between the downstream adaptation and upstream knowledge retention. Empirical analysis indicates that these methods are subject to *parameter space entanglement*, wherein the imposition of downstream adaptation and upstream retention within a unified parameter space results in suboptimal outcomes for both objectives. To address this, we introduce Dual-Pronged LoRA, a novel fine-tuning framework for MLLMs featuring two components. (1) *Task-Decoupled Contextual Routing* (TDCR) makes the first attempt to explicitly decouple tasks by modeling a downstream contextual support region. It adaptively activates the LoRA branch exclusively for downstream-relevant inputs and bypasses it otherwise, thereby preventing downstream adaptation from interfering with pre-trained knowledge. (2) *Downstream-Performance-Oriented Rank Sculpting* (DPO-RS) maximizes downstream adaptation performance by dynamically sculpting ranks via gradient energy. It effectively strengthens critical layers while compressing redundant ones and is theoretically grounded in the minimization of information loss. Extensive experimental results indicate that Dual-Pronged LoRA achieves near-zero forgetting while attaining new state-of-the-art downstream performance.

Large language models spend enormous parameter capacity approximating arithmetic, a deterministic function with an exact solution, and still fail reliably on multi-digit operations. We propose Arithmetic Residual Blocks (ARBs): frozen, differentiable modules inserted into the transformer forward pass that compute exact integer arithmetic using Residue Number System (RNS) encoding on unit circles. The model learns only the interface, a gated injection pathway and a low-rank adapter on the language model head, while the base model and all arithmetic computation remain frozen. The architecture embodies a design principle we call Don't Learn What You Can Compute: any deterministic function expressible as tensor operations can be embedded as a frozen residual block, and the model will learn to route through it because the gradient reward for exact answers dominates internal approximation. We validate this principle by inserting ARBs into a frozen SmolLM2-360M model, training only 1.7M parameters (0.47% of the base model). On exact-match evaluation across addition, subtraction, multiplication, and division with operands up to 3 digits, the augmented model achieves 99.9%+ accuracy across all operations, compared to 4.6% mean accuracy for the unmodified base model. A logit analysis of the frozen base model confirms that the injection pathway overrides the base model's prior uniformly across operations, and that the residual errors (0.047%) concentrate at later digit positions in autoregressive generation, not at any particular operation. These findings motivate the hypothesis that pre-training with ARBs, where no misaligned frozen representations exist, would yield ceiling performance across all operations.

State-of-the-art Extreme Multi-Label Text Classification (XMC) models rely on multi-label attention to focus on key tokens in input text, but learning high-quality attention weights is challenging. We introduce PLANT (Pretrained and Leveraged Attention), a plug-and-play strategy for initializing attention. PLANT works by \emph{planting} label-specific attention using a pretrained Learning-to-Rank model guided by mutual information gain. This architecture-agnostic approach integrates seamlessly with large language model backbones such as Mistral, LLaMA, DeepSeek, and Phi-3. PLANT outperforms state-of-the-art methods across tasks such as ICD coding, legal topic classification, and content recommendation. Gains are especially pronounced in few-shot settings, with substantial improvements on rare labels. Ablation studies confirm that attention initialization is a key driver of these gains. Code and trained models are available at \url{https://github.com/research-anon-487/xcube/tree/plant}.

Zeroth-order (ZO) optimization enables memory-efficient fine-tuning of large language models (LLMs) using only forward passes, but it remains unclear how useful adaptation is distributed across layers. In this work, we reveal a surprising phenomenon: ZO fine-tuning is sharply dominated by a single decoding layer. Across multiple LLM families and downstream tasks, fine-tuning this dominant layer alone consistently matches or even exceeds full-model ZO fine-tuning. We further show that the dominant layer is task-agnostic but model-specific, and can be identified before training through a simple inference-only analysis of activation outliers. Specifically, the dominant layer consistently aligns with the first activation-outlier layer in the pre-trained model. To explain this phenomenon, we analyze how perturbation effects propagate under ZO optimization. We find that the dominant layer combines two key properties: high perturbation sensitivity and early placement in the residual stream, allowing perturbation-induced effects to propagate and accumulate through remaining subsequent decoding layers. As a result, this layer produces disproportionately strong and stable optimization signals under forward-only updates. Extensive experiments on LLaMA2-7B and Qwen3-8B across nine benchmarks show that dominant-layer ZO fine-tuning improves average performance over full-model MeZO and LoRA-based ZO fine-tuning while achieving up to 4.52$\times$ training speedup.

Diffusion-based foundation models for crystal generation have become competent priors over the manifold of plausible materials, but turning that prior into generation of promising candidates requires targeting stability while balancing several conflicting properties at once. Fine-tuning is expensive and produces a new model for every new objective, while differentiable guidance is not applicable to many real-world rewards. We instead adapt these models at \emph{inference time} using Diffusion Tree Search (DTS), a Monte Carlo tree search procedure over the denoising trajectory. We make two technical contributions on top of DTS that are particularly relevant for AI-for-science applications. First, we show how to use existing property-conditional foundation models (e.g., classifier-free guidance variants of MatterGen) as guided proposals inside DTS, with importance-corrected backups and selection that keep the unconditional base model as the target distribution. Second, we extend DTS to multi-objective sampling by maintaining per-objective soft value estimates and selecting expansions through a sampled scalarization, so a single search produces samples spanning a Pareto front while sharing computation across preferences. We instantiate both adaptations on de novo crystal generation: single-objective stability search with different underlying models, conditional generation of crystals belonging to rare space groups, and a multi-objective tasks using a pretrained conditional model. Across these, DTS-based inference-time adaptation improves stability rates and Pareto-front quality without modifying the base models.

Direct diffusion modeling of high-resolution spatiotemporal fields is computationally challenging. Parameter-efficient primitives address this by representing high-dimensional data with a compact set of parameters. In this paper, we construct data-dependent tensor primitives without pretrained compression autoencoders. Our construction starts from Tucker decomposition, which captures low-rank multilinear structure through a core tensor and mode-wise factors. However, Tucker factors are non-unique: the same tensor can be represented by different rotated factors, which complicates generative modeling. We address this issue with orthogonal Procrustes (OP) alignment. Specifically, we select medoid anchor matrices from the data and align the factor matrices to resolve the gauge ambiguity. This yields matrix Grassmannian primitives and tensor Grassmannian primitives that are compact, data-adaptive, and directly decodable by explicit multilinear reconstruction. Theoretically, we prove that the proposed primitive maps are homeomorphisms between low-rank tensors and their corresponding primitive spaces, certifying that the representations are non-degenerate and topologically faithful. Building on these primitives, we propose *Diffusion in Aligned Tensor Space* (DiffATS), a generative framework that trains diffusion models directly on aligned tensor primitives. Across images, videos, and PDE solutions, DiffATS achieves strong unconditional and conditional generation performance while compressing original data by $3.9\times$ to $210\times$, without relying on any pretrained deep compression autoencoders.