NeurIPS 2026 Papers

The rapid growth of large language models is driving organizations to expand their GPU clusters, often with GPUs from multiple vendors. However, current deep learning frameworks lack support for communication across heterogeneous vendor GPUs, leading to inefficiency and higher costs. We present HetCCL, a collective communication library that unifies vendor-specific backends and enables RDMA-based communication across GPUs without requiring driver modifications. HetCCL introduces two novel mechanisms that enable cross-vendor communication while leveraging optimized vendor libraries, NVIDIA NCCL and AMD RCCL. Evaluations on a multi-vendor GPU cluster show that HetCCL matches NCCL and RCCL performance in homogeneous setups while extending support to heterogeneous environments, enabling practical and efficient training across NVIDIA and AMD GPUs without changes to existing deep learning applications.

Deep networks achieve high realism in texture synthesis but often rely on opaque, over-parameterized representations that have limited direct connection to biological vision. Neuroscience-informed pyramid models are interpretable, compact, and biologically grounded, but rely on small, hand-curated statistics sets that are difficult to generalize or extend. We introduce GramStatTexNet, a hybrid analysis-by-synthesis framework that combines the multi-scale Gabor filter structure of classical models with the flexibility of Gramian-based correlations, organized into structured filter families. This structured representation enables analyses that are challenging for deep-feature methods: each statistic carries a named filter-family identity, allowing us to decompose synthesis fidelity into per-family contributions and to learn an interpretable, low-dimensional embedding that preserves perceptual texture content. We compare our framework to classical and deep texture-synthesis baselines across diverse texture categories, and on multiple perceptual metrics, highlighting the compactness and interpretability advantages of this hybrid representation, alongside competitive synthesis quality. This same framework extends naturally to peripheral textures via spatial pooling, and to dynamic textures using a spatiotemporal filter bank. GramStatTexNet provides a unified, interpretable framework for analyzing and modeling visual information with natural extensions across space, time, and eccentricity.

Cross-entropy (CE) is the default training loss for supervised classification, but its sample efficiency is limited when labels are scarce. Existing remedies primarily act on the data side, via augmentation, synthesis, or transfer from pretrained models; the training objective itself is rarely revisited. We revisit it here. Drawing on the classical observation that generative classifiers reach their asymptotic error with fewer samples than discriminative ones, we propose Generative Cross-Entropy (GenCE), a drop-in replacement for CE that introduces a generative learning principle into a standard discriminative network without altering the architecture or fitting a separate density model. GenCE follows from a Bayesian rewrite of the class-conditional likelihood and, in the mini-batch approximation, reduces to normalizing each sample's softmax score against the model's predictions on the batch, coupling the training signal across examples sharing a class. We extend the proper-scoring-rule framework to such non-local losses and prove that GenCE is strictly proper under a mild completeness condition: its population risk is uniquely minimized at the true posterior. Across three datasets, on two architectures and in both balanced small-data and class-imbalanced regimes, GenCE outperforms CE and other widely used losses, while also producing better-calibrated probabilities and stronger out-of-distribution detection.

As large language models (LLMs) scale rapidly, dense full-parameter adaptation becomes increasingly expensive, motivating sparse and modular architectures such as Mixture-of-Experts (MoE) models. This shift raises a key question for parameter-efficient fine-tuning (PEFT): at what granularity should parameters be selected and updated? Existing PEFT methods such as LoRA operate on predefined weight matrices, while expert-level sparse tuning methods update entire selected experts. However, we observe that activated experts are internally sparse, with only a small fraction of intermediate channels strongly responding to downstream tasks, indicating that expert-level adaptation is still too coarse. We propose NSFT (Neural Sub-expert Fine-Tuning), a fine-grained PEFT framework that refines MoE adaptation from experts to sub-experts. NSFT decomposes each expert along the intermediate dimension into structured channel groups and selects task-relevant sub-experts by combining routing importance with intra-expert activation saliency. To optimize sparse partial updates, NSFT further introduces learning-rate scaling and dynamic gradient scaling to compensate for the reduced effective update magnitude. Experiments on OLMoE and Ling-mini-2.0 across challenging domain-specific tasks and general benchmarks show that NSFT consistently outperforms representative PEFT and expert-level sparse tuning baselines, while using substantially fewer trainable parameters and preserving competitive general capability. These results suggest that sub-expert-level adaptation is a more precise and efficient PEFT paradigm for MoE LLMs.

Training-free guidance enables pre-trained diffusion and flow models to optimize application-specific objectives using feedback from external black-box reward functions. However, existing methods are feedback-inefficient because reward feedback is used only transiently to inform a localized gradient approximation or a discrete search decision, and is subsequently discarded. To address this limitation, we propose Flow-Direct, a framework that guides the generation process via a persistent *guidance field*. Theoretically, this guidance field is analytically derived from the log-density ratio between the base and reward-weighted target distributions; it transports the pre-trained distribution to the target distribution. In practice, the field is implemented as a non-parametric estimator constructed from all accumulated reward-evaluated samples. As more samples are collected during optimization, this empirical guidance field becomes increasingly accurate. This persistent formulation yields two major advantages. First, Flow-Direct is highly *feedback-efficient*: because every evaluated sample is used to refine the global guidance field, no reward information is wasted. Second, the framework is naturally *reusable*: once optimization is complete, the collected dataset defines a reusable guidance field for generating novel target samples without additional reward evaluations, and distinct guidance fields can be combined to generate samples that simultaneously satisfy multiple objectives.

Long-context training of large language and multimodal models couples expressive structured attention masks (document, prefix-LM, share-question, sliding-window, block-sparse) with multi-node context parallelism, exposing the mask as a dominant bottleneck. The column-wise sparse mask of FlashMask compresses such masks from $O(n^2)$ to $O(n)$ via four query-axis boundary vectors and covers a broad family of structured patterns. Implemented for Ampere GPUs, FlashMask does not exploit FlashAttention-3's TMA/WGMMA warp-specialized pipeline on Hopper GPUs, and offers no distributed support. We observe that the column-wise sparse mask is indexed by key positions but stores query-axis interval boundaries. Under any multi-chunk context-parallel assignment, this property reduces per-rank mask localization to a provably zero-communication \emph{clip--shift--max} on the boundary vectors, directly compatible with the single-GPU kernel. Building on this property, FlashMask-3 introduces a mask-aware Hopper kernel with a three-role warp specialization atop the FlashAttention-3 TMA/WGMMA pipeline, together with a mask-aware distributed extension: zero-communication mask localization, a sparsity-aware load balancer (HCC-IPO), and a mask-aware compute-communication overlap with mask-driven communication pruning, whose forward all-gather additionally adopts topology-aware hierarchical routing. On a single H100 GPU across 12 structured masks at context lengths from 4K to 128K, FlashMask-3 improves achieved TFLOPs over FlashMask, FlexAttention, and MagiAttention by 45\% to 141\%, 34\% to 71\%, and 2\% to 44\%, respectively. Under context-parallel degrees from 4 to 32 with 8K sequence per device, it attains 31\% to 75\% compute-communication overlap efficiency and improves throughput over Megatron-Core CP (ring), Megatron-Core CP (all-gather) and MagiAttention by $3.25\times$ to $5.11\times$, $1.97\times$ to $3.26\times$ and $1.06\times$ to $4.09\times$ respectively. The source code will be released on GitHub.

k-means has historically been positioned primarily as an offline processing primitive, typically used for dataset organization or embedding preprocessing rather than as a first-class component in online systems. In this work, we revisit this classical algorithm under the lens of modern AI system design and enable k-means as an online primitive. We point out that existing GPU implementations of k-means remain fundamentally bottlenecked by low-level system constraints rather than theoretical algorithmic complexity. Specifically, the assignment stage suffers from a severe IO bottleneck due to the massive explicit materialization of the N × K distance matrix in High Bandwidth Memory (HBM). Simultaneously, the centroid update stage is heavily penalized by hardware-level atomic write contention caused by irregular, scatter-style token aggregations. To bridge this performance gap, we propose Flash-KMeans, an IO-aware and contention-free k-means implementation for modern GPU workloads. Flash-KMeans introduces two core kernel-level innovations: (1) FlashAssign, which fuses distance computation with an online argmin to completely bypass intermediate memory materialization; (2) Sort-Inverse Update, which explicitly constructs an inverse mapping to transform high-contention atomic scatters into high-bandwidth, segment-level localized reductions. Furthermore, we integrate algorithm-system co-designs, including chunked stream overlap and cache-aware compile heuristic, to ensure practical deployability. Extensive evaluations on NVIDIA H200 GPUs demonstrate that Flash-KMeans achieves up to 17.9× end-to-end speedup over best baselines. At the kernel level, FlashAssign and Sort-Inverse Update deliver up to 21.2× and 6.3× speedups. By systematically restructuring execution around underlying hardware constraints, Flash-KMeans delivers mathematically exact, scalable, and highly deployable acceleration across diverse AI workloads.

Building on SODA 2026 work Kearns et al. (2026), we study learning in directed acyclic graphs in which each agent observes only a subset of the input coordinates and may use the predictions of its parents as additional features. In the finite-sample protocol, the message sent along each edge is the agent's prediction vector on a shared sample, so the communication cost is linear in the sample size $m$. In this paper we study this model under communication constraints. First, we revisit the finite-sample analysis of the full-message protocol. We identify a gap in the finite-sample proof of Kearns et al. (2026) and develop a different argument that establishes a finite-sample population-risk guarantee along paths satisfying deterministic $M$-coverage. Second, we introduce a shared-sketch protocol in which each agent sends a $k$-dimensional linear sketch of its prediction vector instead of the full $m$-dimensional message, where $k\ll m$. We analyze this protocol through a reduction to least squares on the compressed sample together with an oblivious subspace embedding argument. The resulting bound separates statistics from communication: the statistical term is still governed by the original sample size $m$, while the additional error due to compression is controlled by the sketch dimension $k$. For Gaussian sketches, we obtain an explicit high-probability population-risk bound, showing that row subsampling is not the only possible communication-statistics tradeoff for networked information aggregation.

Optimal control of partially observable Markov decision processes (POMDPs), in general, requires a complete history. Traditionally, a belief state is used as a sufficient statistic to control POMDPs. However, being defined over a continuous space, its full representation is impossible with a finite-memory controller. This raises the question we address in this work: what must the available finite memory preserve to achieve near-optimal control performance? We study this question through an information-theoretic converse that gives a fundamental lower bound on the loss of any finite-memory controller. The analysis introduces the concept of a decision witness that characterizes belief states whose confusion would lead to unavoidable loss. Motivated by this bound, we propose RECAP, a constructive quantization method to design representative memory states. We prove upper bounds on the optimality gap of any RECAP codebook, and use the resulting components together with the converse to define the codebook objective. Experiments on 11 POMDP benchmarks demonstrate that RECAP improves over finite-memory representation baselines under the same memory budget.

Efficient key-value (KV) cache management is crucial for the practical deployment of large language models (LLMs), yet existing compression techniques often incur a trade-off between performance degradation and computational overhead. We propose a novel gating-based KV cache eviction method for frozen-weight LLMs that achieves high compression ratios with negligible computational cost. Our approach introduces lightweight sink-attention gating modules to identify and retain critical KV pairs, and integrates seamlessly into both the prefill and decoding stages. The proposed gate training algorithm relies on forward passes of an LLM, avoiding expensive backpropagation, while achieving strong task generalization through a task-agnostic reconstruction objective. Extensive experiments across the Qwen2.5-1M, Qwen3, and Gemma3 families show that our method maintains near-lossless performance while evicting up to 70\% of the KV cache. The results are consistent across a wide range of tasks, including long-context understanding, code comprehension, and mathematical reasoning, demonstrating the generality of our approach.

Low-Rank Adaptation (LoRA) is a popular parameter-efficient fine-tuning method for Large Language Models (LLMs), and practical systems usually need to serve many concurrent LoRA adapters to process different tasks. This is challenging due to the inherent tension between \textit{fairness} and \textit{efficiency}, i.e., achieving fairness requires frequent switching among the adapters such that they have similar service quality, while good efficiency requires to minimize adapter switching to reduce memory movement overheads. Existing schedulers can only achieve a fixed trade-off between fairness and efficiency, and their trade-offs are often suboptimal. To tackle this problem, we propose \textbf{F}airness-aware \textbf{E}fficient \textbf{L}oRA \textbf{P}riority \textbf{S}cheduling (FELPS), which explicitly decomposes fairness and efficiency considerations. In particular, for fairness, FELPS adapts the classic proportional fairness and ensures that the amount of service received by each adapter is proportional to its workload. For efficiency, FELPS prioritizes adapters that are already loaded to GPU memory. To conduct scheduling, FELPS combines to the fairness and efficiency terms with an adjustable weight to achieve arbitrary trade-offs between fairness and efficiency. Evaluations on real workloads show that FELPS achieves a superior trade-off between fairness and efficiency than all baselines.

First-person dynamic spatial reasoning requires models to track continuous motion and precise geometric structure, but the quadratic attention cost of Transformer-based Video-LLMs makes dense visual tokens computationally expensive. Existing token pruning paradigms predominantly rely on discrete static snapshots, failing to preserve the physical dynamics essential for reasoning. We propose Event Cascade Pruning (ECP), to our knowledge the first training-free framework that leverages the high-frequency motion cues from event cameras as a continuous event-guided motion prior to guide token selection. ECP combines three stages: Event-Triggered Causal Sampling to anchor motion-informative keyframes, Event-guided Motion Saliency Filtering to remove static or low-motion regions, and Event-Attention Ranking Fusion to calibrate spatial attention with motion salient dynamics. With 80\% visual token reduction, ECP outperforms the full-token baseline (37.62\% vs. 36.31\%) while achieving 1.89× inference speedup and 52\% GFLOPs reduction. We further introduce ESR-Real, the first real-world RGB-event benchmark for first-person spatial reasoning, where ECP improves accuracy by 2.68\% over full-token baselines.

Large language model (LLM) agents typically perform function calling by retrieving relevant tools and reinserting their specifications as textual prompts at inference time. Although retrieval reduces the number of candidate tools, selected tool specifications must still be repeatedly re-encoded, incurring substantial latency and memory overhead, especially in on-device environments. We argue that this inefficiency stems from treating tool knowledge as transient text rather than reusable execution memory. To address this, we propose EngramState, a state-centric framework that compiles tool specifications offline into reusable recurrent state priors and directly loads them at runtime for query-only inference. EngramState combines retention-aware state construction, channel-wise state compression, a lightweight same-backbone state retriever, and feedback-driven state routing. On the DroidCall benchmark with an RWKV-7 1.5B backbone, EngramState reduces prompt tokens from 985 to 22 (97.8\%) under the Top-4 retrieval setting and substantially lowers time-to-first-token (TTFT) on a Galaxy S25 Ultra while maintaining or improving function-calling accuracy. Furthermore, the proposed state retriever achieves competitive retrieval quality using only 8.24\,MiB of retrieval storage. These results demonstrate that reusable recurrent states can serve as an efficient execution-time memory interface for scalable on-device function calling.

Speculative decoding accelerates LLM inference but suffers from performance degradation when target models are fine-tuned for specific domains. A naive solution is to retrain draft models for every target model, which is costly and inefficient. To address this, we introduce a practical parameter- and data-efficient framework named Efficient Draft Adaptation, abbreviated as EDA, for efficiently adapting draft models. EDA introduces three innovations: (1) a decoupled architecture that utilizes shared and private components to model the shared and target-specific output distributions separately, enabling parameter-efficient adaptation by updating only the lightweight private component; (2) a data regeneration strategy that utilizes the fine-tuned target model to regenerate training data, thereby improving the alignment between training and speculative decoding, leading to higher average acceptance length; (3) a sample selection mechanism that prioritizes high-value data for efficient adaptation. Our experiments show that EDA effectively restores speculative performance on fine-tuned models, achieving superior average acceptance lengths with significantly reduced training costs compared to full retraining. Code will be released upon acceptance.

Training Neural ODEs requires backpropagating through an ODE solve. The state-of-the-art backpropagation method is recursive checkpointing that balances recomputation with memory cost. Here, we introduce a class of algebraically reversible ODE solvers that significantly improve upon both the time and memory cost of recursive checkpointing. The reversible solvers presented calculate exact gradients, are high-order and numerically stable -- strictly improving on previous reversible architectures. On scientific modeling and time series classification experiments, reversible solvers reduce the training time by more than $2\times$ while using at least $10\times$ less memory on average than recursive checkpointing.

We study the classical problem of linear regression under heteroskedastic noise, where observations arise from multiple sources with unknown and potentially unequal variances. Specifically, we consider a setting with $n$ sources, each providing a linear measurement of an unknown $d$-dimensional parameter vector $\beta^*$, where each source is associated with an unknown noise variance. We focus on the classic "subset-of-signals'' model, where one assumes that $m$ of the sources have a relatively low noise, formally, with variance of at most 1. Our goal is to accurately estimate $\beta^\*$ despite the presence of high-variance sources. We show that if $m$ is sufficiently large, it is possible to recover $\beta^*$ with sub-constant error. Specifically, we prove a recovery error bound of $\tilde{O} \left(\frac{\sqrt{nd}}{m}\right)$, significantly improving upon previously known results for this setting.

Log-Gaussian Cox Processes (LGCP) are essential for modeling spatial point patterns but suffer from high computational costs and intractable likelihoods. We propose VoGCAM, an efficient Variational Voronoi Gaussian Coordinate Ascent Maximization framework for fitting LGCPs. Our approach first approximates the intractable integral in the LGCP likelihood using a flexible Voronoi tessellation, which incorporates observed points as integration nodes. By applying variational Gaussian approximation, we derive an Evidence Lower Bound (ELBO) that admits an explicit and closed-form expression. To optimize this objective, we develop a novel coordinate ascent algorithm that updates parameter blocks via Newton and fixed-point methods. We further enhance scalability by adopting a Nearest Neighbor Gaussian Process (NNGP) prior and utilizing the Woodbury formula to reduce matrix inversion costs. Theoretically, we prove the existence and uniqueness of the optimal solution and establish the convergence of our algorithm. Numerical experiments on synthetic and real-world data demonstrate that VoGCAM offers superior computational efficiency and inferential accuracy over state-of-the-art methods like INLA and VIFRK.

Spiking neural networks (SNNs) are a promising paradigm for efficient neuromorphic computing, yet their training remains challenging due to the non-differentiability nature of spike firing. Backpropagation through time (BPTT) with surrogate gradients (SG) has become a mainstream approach due to its superior performance. However, this method incurs high computational and memory demands during training, while its update rule depends on specific neuronal dynamics. Furthermore, when targeting neuromorphic hardware implementation, the temporal error-propagation process and the full membrane-potential access required by SG are difficult to realize; these can introduce additional readout overhead, measurement noise, or perturbations to the system dynamics. To address these issues, we propose the input-driven derivative-free training (IDDFT) method. Rather than relying on membrane-potential-based surrogate derivatives, IDDFT constructs derivative-free error signals by applying a relaxed nonlinearity to the inputs, thereby avoiding the temporal error-propagation process and reducing the dependence of the update rule on specific neuronal dynamics. By eliminating access to internal membrane-potential states, our method reduces computational and memory costs, enhances compatibility with neuromorphic hardware, and enables training of black-box SNN models. We validate the IDDFT method through theoretical analysis and systematic experiments, demonstrating its effectiveness as an input-driven, derivative-free mechanism for constructing update signals. When integrated with a biologically inspired training strategy, IDDFT maintains comparable performance while substantially enhancing robustness to the choice of relaxed nonlinearities. Based on the resulting bio-inspired combined framework, we conduct evaluations on CIFAR-10, CIFAR-100, CIFAR10-DVS, and Tiny-ImageNet. Experimental results demonstrate that our method achieves performance comparable to BPTT and state-of-the-art methods, while exhibiting remarkable robustness under adversarial attacks.

LLM-based workflows compose specialized agents to execute complex tasks, and these agents usually share substantial context, allowing KV-Cache reuse to save computation. Existing approaches either manage KV-Cache at agent level and fail to exploit the reuse opportunities within workflows, or manage cache at the workflow level but assume that each workflow calls a static sequence of agents. However, practical workflows are typically dynamic, where the sequence of invoked agents and thus induced cache reuse opportunities depend on the context of each task. To serve such dynamic workflows efficiently, we build a system dubbed PBKV (\textbf{P}rediction-\textbf{B}ased \textbf{KV}-Cache Management). For each workflow, PBKV predicts the agent invocations in several future steps by fusing the guidance from historical workflows and context of the target workflow. Based on the predictions, PBKV estimates the reuse potential of cache entries and keeps the high-potential entries in GPU memory. To be robust to prediction errors, PBKV utilizes the predictions conservatively during both cache eviction and prefetching. Experiments on three workflow benchmarks show that PBKV achieves up to $1.85\times$ speedup over LRU on dynamic workflows, and up to $1.26\times$ speedup over the SOTA baseline KVFlow on the static workflow.