NeurIPS 2026 论文列表

Transformer models rely on attention mechanism to capture long-range dependencies but suffer from quadratic complexity, limiting their scalability to long sequences. Kernel-based linear attention reduces this complexity but typically relies on fixed or weakly learnable kernels, restricting expressiveness and performance. In this work, we propose Flexformer, a flexible linear Transformer that learns attention kernels in a fully data-driven manner. Flexformer builds on random Fourier feature-based linear attention and treats spectral frequencies as trainable parameters, enabling the model to learn a broad family of attention kernels. We develop both stationary and nonstationary variants, with the latter offering strictly greater expressiveness. Extensive experiments on language modeling and sequence classification demonstrate that Flexformer consistently outperforms baselines. Moreover, Flexformer can be effectively distilled from pretrained Transformers to recover softmax attention and exhibits strong kernel transferability across domains, achieving both high efficiency and competitive performance on long-sequence tasks.

A cover song re-renders an existing piece by preserving its tonal content—melody and chord progression—while reshaping other musical attributes. This suggests a natural formulation for generative cover synthesis with pretrained music language models: explicitly control tonal content, while specifying lyrics and style through text. Existing approaches condition on frame-level features, tightly coupling generation to the reference’s timing and structure, and limiting flexibility such as tempo variation, structural rearrangement, or partial conditioning. To this end, we present FlexCover, a cover generation model that conditions a pretrained text-to-song foundation model on a symbolic lead sheet. Our design enables alignment-free control: generated outputs preserve the tonal signature of the source without frame-level alignment, allowing flexible timing, structure, and segment-level generation. We further introduce a training curriculum with partially mismatched audio–symbolic pairs, improving diversity and robustness. We evaluate FlexCover with objective metrics, standard subjective ratings, and in-depth expert interviews that provide fine-grained diagnostic insights. Experiments show that FlexCover achieves state-of-the-art performance among open-source systems and is competitive with leading commercial models such as Suno-v5.5.

Recent fMRI foundation models differ substantially in the spatial scale at which they represent brain activity. ROI- and connectivity-based models are efficient but coarse, whereas voxel-level models preserve fine-grained spatial structure but require specialized architectures and costly fMRI-specific pretraining. We ask whether part of this performance gap reflects the importance of preserving cortical geometry. Motivated by evidence that macroscale brain activity is strongly constrained by brain geometry, we introduce FlatClip, a training-free surface-level baseline that renders cortical activity as geometry-aware flatmap sequences and reuses a frozen SigLIP2 image encoder with only a lightweight downstream probe. Because FlatClip is not pretrained on fMRI, it provides an out-of-domain reference for evaluating how much information can be recovered from cortical geometry without relying on fMRI-specific pretraining data. Across resting-state and visual-fMRI benchmarks, FlatClip outperforms ROI/FC-style controls and several general-purpose fMRI foundation baselines, while voxel-level models remain a strong upper comparison. Geometry-control experiments show that disrupting cortical topology reduces performance, whereas mapping ROI-level signals back into atlas-defined 4D volumes partially recovers performance but remains below full voxel inputs. Together, these results position cortical geometry as a key factor in fMRI representation learning and establish surface-level flatmap sequences as a practical middle-ground baseline between ROI and voxel models.

3D geometric understanding is undergoing a major transition where rigid, deterministic, optimization-based pipelines are being replaced by large-scale, deep learning models. While recent methods achieve impressive performance on tasks such as multi-view 3D reconstruction and camera pose estimation, they suffer from two critical limitations: $\textbf{(i):}$ current state-of-the-art (SOTA) models require industry-level compute and training corpora, severely limiting accessibility, and $\textbf{(ii):}$ recent approaches remain constrained to modifying existing pretrained transformer models, since substantially changing their architecture would require retraining from scratch. We instead study a complementary direction: training an attention-free state space model for dense 3D understanding. We introduce $\textbf{FlashSSM-3D}$, the first 3D reconstruction state space model capable of performing accurate geometric reasoning over hundreds of images with linear computational complexity. However, training such a model from scratch with competitive performance is impractical under typical academic compute constraints. To address this, we train FlashSSM-3D by distilling knowledge from VGGT, avoiding the need to learn 3D geometric priors entirely from scratch. However, standard distillation losses are poorly suited for dense 3D reconstruction, where preserving real-world spatial structure is essential. We therefore propose a $\textbf{geometry-aware knowledge distillation}$ method specifically tailored for large-scale 3D reconstruction models. Our method introduces two novel losses: a Procrustes-based alignment loss that enforces global structural consistency between student and teacher 3D point predictions, and a camera reprojection loss that directly supervises view-consistent camera geometry between teacher and student. Extensive experiments across multi-view 3D reconstruction and camera pose estimation benchmarks demonstrate that FlashSSM-3D achieves competitive performance while requiring only $\sim25\\%$ of the training data and less than $10\\%$ of the training compute used by VGGT. At inference time, FlashSSM-3D benefits from the practical efficiency of state space models, reconstructing scenes from 500+ views in under 10 seconds while maintaining competitive accuracy.

Portrait video relighting aims to modify the illumination of a portrait video while preserving photorealism and temporal stability. The task becomes particularly challenging under dynamic lighting, where illumination direction, intensity, and color vary over time. To tackle this problem, we introduce FlashRelight, a novel two-stage framework for dynamic portrait video relighting. FlashRelight first suppresses the source illumination through delighting, and then synthesizes the target dynamic lighting on an illumination-neutral portrait video. FlashRelight supports three complementary control modalities: envmap sequences for explicit frame-level lighting control, dynamic lighting text for semantic illumination editing, and lighting reference videos for visual lighting transfer. Training follows a low-cost hybrid strategy that combines static one-light-at-a-time (OLAT) data with generated dynamic-lighting portrait videos, covering both controlled illumination variation and realistic portrait motion. Extensive experiments show that FlashRelight achieves temporally coherent and photorealistic dynamic relighting while enabling flexible cross-video lighting transfer.

Diffusion and flow matching have emerged as expressive generative planners for autonomous driving planning, owing to their ability to model high-fidelity and multi-modal trajectory distributions. Nevertheless, existing generative planners are predominantly optimized through imitation learning, which induces a fundamental mismatch between supervised trajectory fitting and closed-loop planning metrics, while their iterative sampling procedures often impose substantial computational overhead. In this paper, we propose FlashPlanner, a goal-conditioned flow-matching planner with online RL finetuning for closed-loop AD planning. FlashPlanner introduces a continuous future goal point as a compact navigation interface for the generative planning policy. This goal is produced by a continuous goal predictor, which is first pretrained and subsequently optimized with RL using closed-loop feedback from multiple candidate goals at each decision step. To support efficient online reinforcement fine-tuning and real-time deployment, FlashPlanner adopts a data-prediction flow-matching objective and removes redundant architectural components in existing diffusion-based planners. Experiments on the closed-loop nuPlan and interPlan benchmarks demonstrate that \textit{FlashPlanner} achieves state-of-the-art planning performance while delivering 6× faster inference (166 FPS) than the previous SOTA baseline (28 FPS). We will open-source our project.

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.

The Feed-Forward Network (FFN) dominates the activation memory of modern Transformers because the intermediate tensor of width d_ff is materialized in HBM between the two matmul stages. FlashAttention-style I/O-aware tiling avoids the analogous cost in attention, yet has not been brought to FFN. We identify the reason as a hardware-geometry constraint: under one-pass tiled accumulation a fused FFN kernel must keep an output accumulator of width d_model in on-chip SRAM, which is infeasible whenever d_model >= 2048 on currently shipping datacenter GPUs. We formalize this as an SRAM-feasibility proposition and prove a corollary: multi-head decomposition is the simplest architectural transform that brings FFN inside this SRAM-feasible regime, because it makes the per-head accumulator width d_h = d_model/H a free parameter that can be driven below the GPU's SRAM cutoff. Naive multi-head FFN, however, blows up the per-head expansion ratio d_ff/d_h as d_model grows, drifting away from the well-validated SwiGLU 8/3 optimum and degrading quality at scale; we resolve this with a dense MoE-like sub-network correction that restores the optimal ratio while preserving SRAM-feasibility. The resulting architecture, FlashFFN, is an SRAM-feasible FFN block that replaces SwiGLU under matched parameter and FLOPs budgets and derives directly from the SRAM constraint. Across 128M--1.3B parameter models trained on 60--100B tokens of The Pile, FlashFFN achieves a Pareto improvement over SwiGLU at fixed compute: up to ~4x lower peak activation memory at long sequences, inference latency that matches or improves over SwiGLU (up to ~1.32x faster on H100), and consistent quality gains (up to -0.84 in evaluation perplexity and +1.60 average accuracy across six downstream benchmarks at 1.3B)

LLM-based evolution has emerged as a promising way to improve agents by refining non-parametric artifacts, but its wall-clock cost remains a major bottleneck. We identify that this cost comes from synchronized stage execution and imbalance inside each LLM-heavy stage. We present FlashEvolve, an efficient framework that replaces synchronized execution with asynchronous workers and queues, allowing different stages and steps to overlap. To handle data staleness introduced by asynchrony, FlashEvolve tracks artifact versions and applies different policies to update, discard, or patch stale artifacts. Unlike weight-space staleness in asynchronous RL, language-space staleness is inspectable and repairable: a stale artifact is not just delayed work, but readable evidence that the LLM can reflect on, revise, and turn into useful evolution signal. FlashEvolve further improves throughput and token efficiency with speculative stage completion and adaptive workflow control. On GEPA workloads, FlashEvolve improves proposal throughput by 3.5x on local vLLM and 4.9x on API serving over synchronous GEPA. The same design also applies to ACE and Meta-Harness. Our code repository is available at https://anonymous.4open.science/r/FlashEvolve-FEC7/.

Diffusion models have become a leading approach in visual generative modeling, particularly effective in text-to-image and controllable generation. Their main drawback is slow inference as it requires many sampling steps. While recent distillation methods enable one-step text-to-image generation, extending these techniques to controllable generation is still challenging. Existing controllable models, such as ControlNet, rely on dual-branch design that make distillation difficult. In this work, we adopt a simpler strategy. Instead of adding extra conditional branches, we finetune a selected subset of layers in the original text-to-image diffusion model to directly enable controllable generation. This results in DeltaControl, a single-flow, multi-step diffusion model capable of supporting spatial-condition generation while matching the performance of dual-branch approaches. Building on this teacher model, we further apply step-distillation to obtain FlashControl, a one-step model capable of spatial controllable generation with a single forward pass.

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.

Variance reduction techniques improve the efficiency of drawing conclusions from data. Recent proposals that incorporate machine learning predictions into estimators based on weighted sampling use the Poisson Sampling (PS) algorithm, which independently flips a biased coin at each unit and includes all heads in the sample. A fundamental flaw with PS is that it has a random sample-size, which presents practical challenges and introduces variance into the estimator. We show that a fixed-size alternative, known in survey-sampling as Conditional Poisson Sampling (CPS), is a drop-in replacement for PS on general multivariate Z-estimation problems, which includes convex M-estimation considered in prior work. Notably, we show that the asymptotic variance of CPS is no worse than that of PS, and can be substantially better in realistic scenarios where predictions are imperfect. We illustrate these advantages on both synthetic and real datasets. In light of our contributions, there is no need to suffer a variable sample-size.

Looped-in-depth architectures provide an inductive bias toward learning step-by-step procedures for tasks that require compositional reasoning. The effective depth reached by looping determines the quality of the solution. Similar to deep architectures, looped architectures are prone to signal propagation issues as the halting decision is postponed. In this paper, we address these signal propagation issues by using pre-norm layers and residual scaling. Furthermore, we propose ***FPRM***: a Fixed-Point Reasoning Model that uses fixed-point convergence as an end-to-end halting mechanism in a looped architecture. We show that fixed-point halting allows FPRM to adapt its compute to the difficulty of the task. FPRM proves effective on common reasoning benchmarks, namely sudoku, maze, and state tracking.

Masked Generative Models (MGMs) enable parallel decoding and achieve strong performance across modalities, but require full-sequence bidirectional transformers at every step, making training costly and degrading quality under low sampling budgets. Existing work improves efficiency via better samplers or fixed-depth backbones, but does not vary the depth of the denoiser across sampling steps. We introduce Fixed-Point Masked Generative Models (FP-MGMs), which replace part of the denoiser with a fixed-point solver over shared attention layers to enable adaptive depth with fewer parameters. To make it more effective for masked generation, we first introduce a cross-step consistency loss, which aligns hidden representations at neighboring denoising steps and, second, three-state reuse (3SR) which warm-starts the solver using the previous solution by treating unchanged, still-masked, and newly revealed tokens differently. Together, these components define our complete training-to-inference framework for fixed-point masked generation, \emph{CoFRe}. We also show that pre-trained MGMs can be converted into FP-MGMs with short fine-tuning, avoiding full retraining. Across modalities, CoFRe improves the quality and cost trade-off. On OpenWebText, CoFRe reduces parameters by 38.8\%, training time by 11.5\%, and VRAM by 16.9\%, while improving generative perplexity from 830.8 to 101.8 at a budget of $96$ transformer-block forward passes, compared to MDLM. In ImageNette, CoFRe reduces training time by 48.6\% and VRAM by 50.7\%, while improving FID in all sample budgets tested. Overall, CoFRe offers a practical framework for cheaper training and stronger low-budget masked generation.

We introduce \emph{universal transformers}: fixed-parameter architectures that can simulate any transformer in a given class via a suitable input embedding. Analogous to a universal Turing machine, the input embedding encodes a description of the target model while all internal parameters remain fixed. We provide explicit sparse constructions achieving universality when the embedding dimension is sufficiently large, and further show that universality is generic: randomly initialized transformers are universal almost surely, which aligns with recent empirical results of Zhong and Andreas (2024). We empirically validate our theory on the algorithmic tasks of parenthesis balancing and multi-hop reasoning. Our results suggest that much of a transformer’s expressive power may reside in its input representation rather than its learned weights.

Long-context LLM agents often struggle with growing token, memory, and latency costs, making efficient context compression essential for practical deployment. Existing LLM-as-a-compressor methods remain noticeably inferior to using the full context. We find that this gap partly stems from their inability to preserve contextual information effectively. In this work, we revisit context compression from a structural perspective and identify two key bottlenecks in standard LLM-based compressors: limited coordination among compression tokens during information aggregation, and layerwise dilution that weakens useful signals from intermediate hidden states. To address these limitations, we propose **ComprExIT** (Context **Compr**ession via **Ex**plicit **I**nformation **T**ransmission), a new context compression framework based on explicit information transmission. ComprExIT adaptively selects features across frozen LLM layers, then allocates information from anchors to compression slots through a globally coordinated transport plan. Experiments on 12 datasets show that ComprExIT consistently outperforms strong soft-compression baselines, improving average F1 by up to 18.5%, while adding only ~1% trainable parameters and achieving more than 2× faster compression than the fastest baselines. The code will be released upon acceptance.

Reliable hypothesis testing is the foundation of many empirical scientific claims. Large language model (LLM) agents are increasingly used to automate this process, as they can inspect datasets, generate code, and produce analyses end-to-end. However, we show that they frequently make subtle inferential errors that lead to incorrect conclusions despite correctly executed analyses. Existing benchmarks fail to capture this failure mode, as they rarely assess whether a reported p-value is statistically valid given the assumptions underlying the data. We address this gap by building P-Bench, a benchmark comprising 425 open-ended, realistic hypothesis-testing tasks spanning economics, biology, and medicine. Each task requires an agent to select a statistical method, compute a p-value, and draw a conclusion given only a scientific hypothesis and a dataset. We further introduce Fisher-R1, an open-weight LLM agent trained for rigorous hypothesis testing using synthetic tasks and reinforcement learning. On P-Bench, Fisher-R1-14B substantially improves over its backbone and outperforms strong proprietary and open-source baselines, including GPT-5.4 and DeepSeek-V4-Pro, achieving a 21% average relative improvement in single-trial success over DeepSeek-V4-Pro, with gains up to 26% on the most challenging tasks. Our results demonstrate that current LLM agents lack reliable statistical reasoning for hypothesis testing and that reinforcement learning on tasks with verified statistical reward substantially improves reliability.

Modern models can have high average accuracy and benign loss tails while remaining fragile on rare subpopulations. We identify an information-geometric mechanism: after nuisance projection, the task direction can become locally nonidentifiable even when loss and raw Fisher look well behaved. Fisher-Glass certifies this failure by applying CVaR to inverse nuisance-projected task information. We prove a loss--information separation theorem and show that, under a boundary-mass condition, the hidden-environment tail of inverse projected Fisher controls robust classification sample complexity. We then derive stable ridge certificates and repair principles: reserves lift Fisher-null tails, portfolios increase collapse codimension, and tail-transverse influence selects counterfactual repairs by combining Fisher-Glass directionality with loss leverage. Synthetic experiments validate the theory; WILDS/Waterbirds diagnostics show that weakest identifiability need not mean lowest accuracy; and Waterbirds training plus counterfactual repair show that Fisher-Glass controls an identifiability axis complementary to loss.

A fruitful approach towards understanding generalization in machine learning involves characterizing how models separate generalizable structure from idiosyncratic or even corrupted features of training data. Fisher information has emerged as a diagnostic tool for this separation, leveraging loss geometry to distinguish generalization from memorization. Here we investigate this connection in vision models trained on CIFAR-10 with controlled label noise that can be learned only through memorization. We show that projecting model weights onto the top Fisher eigenvectors decouples general task performance from noisy sample memorization, and we identify the layer in which this decoupling emerges. We show that the same phenomena are present in small transformers trained to perform in-context learning and reproduced in closed form in a simple, analytically solvable linear model. Our results reveal that Fisher information captures aspects of the geometric structure underlying how neural networks allocate capacity between generalization and memorization.