NeurIPS 2026 Papers

Diffusion models have shown promise in learning to solve constraint optimization problems. However, they are mostly restricted to problems with binary variables and rely on graph neural networks, hindering their application to a broader range of problems such as those with general discrete variables or constraint structures that necessitate global rather than local reasoning. We investigate the use of Diffusion Transformers to address the aforementioned limitations. A naive implementation performs poorly due to a fundamental mismatch between the standard diffusion process and constraint solving: while the former applies *small, incremental denoising* across *all* variables, the latter requires *substantially altering specific subsets of variables* to attain feasibility or optimality. Our method, **Blo**cked **G**ibbs **Di**ffusion **T**ransformer (BloGDiT), is the first to address this limitation by replacing standard joint Gaussian denoising with Blocked Gaussian denoising. BloGDiT uses iterative block resampling and anneals the block size over time to facilitate large, targeted edits within a block of variables. Across Sudoku, Graph Coloring, Maximum Independent Set, and MaxCut, BloGDiT matches or outperforms existing methods, demonstrating that blocked Gibbs-style diffusion provides a highly effective inductive bias for Transformer-based constraint satisfaction and optimization.

We study long-context balanced entropic optimal transport (OT) attention on TPU hardware through a stopped-base, fixed-depth tail-refinement surrogate. After a stopped $T$-step Sinkhorn solve, we unroll a short refinement tail and differentiate that surrogate exactly. For the production $R=2$ case, the backward pass contains four staircase plan factors. We prove an exact one-reference-tile schedule: the $R=2$ score cotangent is a single reference plan tile times an explicit modifier field built from vector cotangents and dual differences. This yields block-wise cost $O((T+R)LW)$, $O(Ld)$ input storage, and $O(L)$ additional HBM usage for fixed head dimension $d$ and band width $W$. We also formalize the current \texttt{dustbin\_block} path as the same balanced surrogate on an augmented support, so the schedule lifts to the gap-aware transport path used in our TPU runs. We provide a local surrogate-bias bound, an a posteriori bias certificate, and a projective contraction certificate for strictly positive active blocks. On synthetic masked problems, the optimized kernel matches exact autodiff of the same centered surrogate to within $10^{-5}$--$10^{-10}$. On TPU v6e-8, a four-configuration Pfam screen completes end-to-end, and a promoted balanced $R=2$ run sustains roughly $8.5$ examples per second through a three-hour budget, reaching step $1437$. Held-out Pfam test shards improve reconstruction from $3.17$ to $0.99$ and sparse CE from $5.86$ to $5.69$ relative to step $0$. These results support exact fixed-depth backward theory, a theorem-matching gap-aware bridge, and trainability evidence for the production path.

Recently, reinforcement learning (RL) has been widely applied during post-training for diffusion large language models (dLLMs) to enhance reasoning with block-wise semi-autoregressive generation. Block size has therefore become a vital factor in dLLMs, since it determines the parallel decoding granularity and affects the rollout trajectories during RL optimisation, e.g., GRPO. Instead of investigating the effect of block size during inference on individual domains, this paper studies block size from a domain conflict perspective for dLLM RL post-training in multi-domain scenarios. The main contributions are: (1) a formulation of domain block size conflict in multi-domain RL for dLLMs, which will largely affect the post-training effectiveness for rollout-based RL methods; (2) a novel dataset, \textbf{Block-R1-41K} is constructed with a best-improved training block size for each sample, which also induces a Block Size Conflict Score to quantitatively measure the domain conflict; (3) a new benchmark, Block-R1, for flexible RL post-training for dLLMs in both single and cross domain; and (4) a simple yet powerful cross-domain post-training method with sample-level best-improved training block sizes. Extensive experiments on 13 distinct datasets, 7 latest RL algorithms, and various different dLLM backbones are covered in Block-R1. The benchmark is open-sourced at https://anonymous.4open.science/r/Block-R1-2026, with the dataset released at https://huggingface.co/datasets/dLLM-R1/Block-R1-41K.

Post-training pruning of large language models (LLMs) zeroes a fraction of the weights and reconstructs the survivors without retraining, addressing the weight-memory bottleneck of LLM serving. At a fixed mask the per-row reconstruction admits a closed-form Optimal Brain Surgeon (OBS) optimum, but the cost of the exact joint solve is cubic in the layer dimension; one-pass methods such as SparseGPT and Wanda approximate it at quadratic cost without attaining the per-block joint optimum. We introduce Block-OBS-GS, which solves the per-block joint OBS reconstruction exactly within SparseGPT's complexity class via a per-row Cholesky factorisation of the in-block damped Hessian (one third the cost of an explicit inverse, reused across the algorithm) and follows it with a single cross-block Gauss--Seidel sweep. We prove three guarantees: per-layer FLOP count in SparseGPT's complexity class, with the leading $mn^{2}$ constant dropping from $1/2$ (SparseGPT) to $3\rho/2$ (Block-OBS-GS) where $\rho=1-s$, hence strictly lower for sparsity $s>2/3$; per-row reconstruction error at most SparseGPT's, strict whenever the activation Hessian has non-zero off-diagonal coupling; and a closed-form distance bound to the unregularised joint-OBS optimum that contracts geometrically per sweep. Across six Llama-3 and Qwen-3 backbones from 3B to 70B parameters, Block-OBS-GS attains the lowest WikiText-2 perplexity at every tested sparsity on the three Llama-3 backbones, with a $47\%$ perplexity reduction over SparseGPT on Llama-3.1-8B at $s=0.875$; the highest seven-task commonsense Avg7 on six of twelve populated cells; 5-shot Massive Multitask Language Understanding (MMLU) within $0.026$ of SparseGPT across populated cells; and the lowest WikiText-2 perplexity on four of five backbones under $2{:}4$ structured sparsity.

Vector quantization is a fundamental primitive for scalable machine learning systems, enabling memory-efficient storage, fast retrieval, and compressed inference. Recent rotation-based quantizers such as EDEN, RabitQ, and TurboQuant have introduced strong guarantees and empirical performance, but the surrounding comparisons have been difficult to interpret because they rely on different distortion criteria, probability regimes, and implementation assumptions. As our first contribution, we provide a unified theoretical comparison of these methods and show that their relative advantages are criterion-dependent rather than absolute: TurboQuant is favorable for MSE distortion, EDEN is effective for expected inner-product distortion, and RabitQ provides strong high-probability control. This comparison clarifies the design principles behind each advantage and shows that no existing method uniformly dominates across the relevant measures. As our second contribution, we introduce *Block-Sphere Quantization* (**BlockQuant**), a new rotation-based block quantization algorithm designed around the spherical geometry of randomly rotated vectors. Unlike coordinate-wise quantizers, **BlockQuant** quantizes blocks on the sphere, preserving the geometry of rotated embeddings more faithfully. We prove that this block-spherical design improves expected distortion performance, including both reconstruction MSE and expected inner-product distortion. Experiments on real embedding data support these theoretical improvements.

Modern large language models increasingly require long contexts for reasoning and multi-document tasks, but attention's quadratic complexity creates a severe computational bottleneck. We present *Block Sparse Flash Attention* (BSFA), a drop-in replacement that accelerates long-context inference while preserving model quality. Unlike methods that predict importance before computing scores, BSFA computes exact query-key similarities to select the top-$k$ most important value blocks for each query. By comparing per-block maximum scores against calibrated thresholds, we skip approximately 50% of the computation and memory transfers for pruned blocks. Our training-free approach requires only a one-time threshold calibration on a small dataset to learn the per-layer and per-head attention score distributions. We provide a CUDA kernel implementation that can be used as a drop-in replacement for FlashAttention. On Llama-3.1-8B, BSFA achieves up to $1.13\times$ end-to-end speedup on LongBench with only a $1.1\%$ accuracy drop, and up to $1.24\times$ on Needle-in-a-Haystack retrieval at a $1\%$ accuracy drop. The attention kernel itself accelerates by up to $1.38\times$. We compare BSFA against five recent sparse attention baselines (SpargeAttention, MInference, FlexPrefill, XAttention, and BLASST), and verify the method on Qwen2.5-7B and on A6000 and H100 GPUs. The implementation is available at https://github.com/Anonymous44414/Block-Sparse-Flash-Attention.

We study an endogenous nonstationary stochastic bandit problem with latent linear dynamics, where actions affect both immediate rewards and the future evolution of an unobserved latent state. Rewards are bilinear in the current action and latent state, inducing history-dependent rewards and a nontrivial long-horizon planning problem. The existing explore-then-commit approach (Choi et al., 2026) achieves $\tilde{\mathcal{O}}(T^{2/3})$ regret by uniformly exploring to estimate the latent dynamics and then committing to an optimized open-loop action sequence. We show that this rate can be improved via adaptive block-level optimism. Our key step is a cyclic approximation: under stable dynamics, the infinite-memory reward process can be truncated, and the open-loop benchmark can be approximated by optimizing a finite-memory block-level proxy. Building on this reduction, we propose a UCB-based block algorithm that maintains confidence sets for the truncated dynamics parameters and selects blocks optimistically. We prove a regret bound of order $\tilde{\mathcal{O}}(\sqrt T)$, significantly improving over the previous $\tilde{\mathcal{O}}(T^{2/3})$ guarantee for the same model. To the best of our knowledge, this is the first $\tilde{\mathcal{O}}(\sqrt T)$ regret guarantee for latent linear-dynamics bandits with bilinear reward observations and an open-loop action-sequence benchmark.

Multimodal Large Language Models (MLLMs) have shown promise as interactive agents, yet precise, programmatic 3D graphics editing remains difficult to train because existing systems largely rely on closed-source model APIs and lack deterministic, state-verifiable execution environments. We present BlenderFORGE, a trainable data-environment-optimization framework for open-weight MLLM agents that edit explicit Blender scenes through Blender Python (bpy) code generation. BlenderFORGE targets four fundamental state-verifiable editing task families: object placement, shape-key editing, lighting adjustment, and material modification. It integrates three components: (1) a scalable perturbation-based data generation pipeline built from curated assets and large-scale 3D scene repositories; (2) an isolated Blender Sandbox for deterministic code execution, structured state export, and closed-loop reward computation; and (3) a two-stage training pipeline that uses offline teacher trajectories for cold-start Supervised Fine-Tuning (SFT) and sandbox-computed rewards for multi-turn tool-agentic Reinforcement Learning (AgentLoop-RL). Experiments across the four task families show that the resulting open-weight Qwen3-VL-8B-BlenderFORGE agent substantially improves over its base model and achieves competitive performance with several zero-shot proprietary MLLM baselines.

AI weather models are fast and skillful, yet their errors are strongly conditional because the best member changes with lead time, variable, initialization state, and atmospheric regime. Operational deterministic multi-model forecasting therefore requires more than averaging strong forecasters. It requires anticipating which member deserves trust before future analyses exist. BlendCast turns this anticipation problem into a vision-language forecasting task. Given meteorological maps and structured ensemble diagnostics, the model reads the current weather situation, reasons about prospective member reliability, and converts that judgment into deterministic per-variable ensemble weights. To teach this behavior, BlendCast first gives the model demonstrations of what member-skill reasoning and weight decisions should look like, then refines the resulting policy with multi-variable reinforcement learning so that its own analyses are rewarded when they lead to better, physically consistent forecasts. On held-out 2025 forecasts, BlendCast reduces WRMSE by 2.6\% over equal weighting and by 2.0% over conventional ensemble post-processing, closing 63.2% of the gap to oracle blending. These results show that, with multimodal weather evidence and reward-aligned training, VLMs can act as real-time forecasters of forecasters for deterministic AI weather ensembles. Project and demo:https://anonymous.4open.science/status/BlendCast-2567

Access to modern generative systems is often restricted to querying an API (the ``black-box" setting) and many properties of the system are unknown to the user at inference time. While recent work has shown that low-dimensional representations of models based on the relationship between their embedded responses to a set of queries are useful for inferring model-level properties, the quality of these representations is highly sensitive to the query set. We introduce the discriminative factorization to distinguish between high- and low-quality query sets in the context of black-box model-level classification. Under this framework, the probability of chance-level classification decays exponentially in the query budget. On three auditing tasks, estimated factorization parameters predict the empirical performance decay rate. We conclude by showing that query sets selected using the estimated discriminative field reproduce the empirical ordering of oracle query sets.

In practice, most commercial LLM providers do not publicly release details of underlying LLM architectures. However, prior work has shown that given limited API access to an LLM (namely, top-$k$ logits and/or a logit bias function), one can recover certain architectural details of an LLM, such as the hidden dimension of the feed-forward network. Perhaps in response to these results, most commercial LLM providers have restricted their APIs to expose only the single logit for each decoded token, and they no longer give users the ability to bias logits. We show that even under current restrictive APIs, several architectural parameters are still recoverable. We present NightVision, an attack that uses restrictive black-box API access to estimate the hidden dimension, depth, and parameter count of an LLM. Algorithmically, NightVision relies on a novel common set prompting technique in which multiple prompts expose log probabilities for the same set of output tokens; a spectral analysis of these results is used to infer hidden dimension. NightVision additionally uses end-to-end time to first token (TTFT) measurements and the estimated hidden dimension to estimate depth and parameter count. We empirically evaluate NightVision on 32 open-source LLMs, recovering hidden dimension to within 23\% average relative error across all models (9\% on MoE models), and depth and parameter count to within 53\% for models exceeding three billion parameters. We run extensive ablations to demonstrate how these accuracies scale with token budget and model properties. Overall, our results suggest that current LLM APIs are not sufficiently restricted to fully obfuscate the architectural details of their underlying models.

We study mathematical programs with equilibrium constraints, in which a leader knows their own cost function, but lacks a model of the followers' response. Instead, the leader can only query this response at specific points. While this setting precludes the use of gradient-based methods, existing zeroth-order approaches treat the composed objective \emph{entirely} as a black box, deploying zeroth-order tools across both the leader and follower. Such approaches are inefficient, as they discard information the leader already possesses about their own cost function. In this work we instead propose to deploy zeroth-order tools only where they are truly needed: to handle the unknown, non-smooth followers' response. Specifically, we first propose \algpzos, an algorithm that combines exact partial gradients of the leader's cost with zeroth-order Jacobian estimates of the followers' response in a chain-rule-inspired manner, and establish that it achieves a strictly lower variance bound than the black-box baseline. Second, we introduce the partial Goldstein subdifferential, a stationarity notion tailored to this composite structure, and prove convergence of our algorithm to both standard and partial Goldstein stationary points. Finally, we validate our method on two application domains -- toll optimization in routing games and defense-attack investment in security games -- demonstrating consistent improvements over black-box baselines in convergence speed, objective value, and estimator variance, with robust performance even under few queries per iteration.

A language model agent acts on a context assembled by retrieval, summarisation, prompt templates, and memory. How many decision-relevant bits must this context carry before the agent can act with low regret? We give an answer that does not depend on the model, decoder, or scaffold. Modelling the agent as a Markov chain H→C→A and treating Bayes regret (the value gap between the chosen and the optimal action) as the distortion measure, we define a regret rate-distortion (RRD) function Rℓ(ε) whose converse is a one-line consequence of the data-processing inequality (DPI). A context that carries fewer than Rℓ(ε) bits cannot drive expected regret below ε, regardless of decoder, scaffold, scale, or chain-of-thought. For finite M-ary decisions under zero-one loss the bound reduces to the Fano floor in closed form, requiring ≈3.14 bits at (M=16, ε=0.1) and ≈6.73 bits at (M=256, ε=0.1). We instantiate the theory on 441 controlled single-step tool-selection cells, a balanced 16-way classification proxy spanning three open-source large language models (Qwen3-32B, Gemma3-27B, GLM-4-32B) and three public benchmarks (API-BANK, τ-bench, ToolBench). Because the gold tool index X=f(H) is a deterministic function of the history, the action-side rate IMM(X; A) lower-bounds the context capacity by data processing, so an action-side test of the converse is strictly tighter than a context-side one. Every cell respects the predicted floor (median margin +1.07 bits). Across 629 matched-rate pairs, 97.5% share overlapping regret confidence intervals (CIs), and a chain-of-thought supplement slides cells along the frontier rather than across it. The same data exposes a tokens-versus-bits gap reaching ∼1.7×10⁶:1, with the model extracting fewer than 0.05 bits about the gold action on a 32,768-token ToolBench prompt. In the single-step regime, the framework supplies a falsifiable, model-agnostic stopping criterion at IMM=Rℓ(ε) and reframes context engineering as an information-allocation problem.

The practical deployment of large language models on resource-limited devices is constrained by three coupled costs: memory footprint, energy consumption, and inference latency. Ultra-low-bit backbones such as BitNet b1.58 mitigate the first two issues by ternarizing the weights to 1.58 bits, but the inference latency remains bounded by autoregressive next-token prediction (NTP), which commits only a single token per forward pass. Multi-token prediction (MTP) improves inference latency by committing several tokens per step. However, the naive combination of MTP with a 1.58-bit backbone does not yield end-to-end gains: once 1.58-bit execution reduces the arithmetic cost, the runtime cost of the accepted-prefix dynamics of MTP dominates each iteration and erodes the savings. In this work, we propose BitMTP, an algorithm--system co-design framework that makes MTP native to a single 1.58-bit backbone. BitMTP is organized around one design principle that we term accepted-prefix alignment: at every iteration, the committed prefix is simultaneously a valid autoregressive continuation, a directly promotable KV slice, and the index of a pre-captured replay state. Three runtime mechanisms enforce this principle along the kernel, memory, and scheduling axes, respectively: a role-width-specialized 1.58-bit execution primitive, a finite-state graph atlas with shared-KV memory over a geometric segment ladder, and an optimistic verify-prefetch scheduler on the verifier-bearing decoding branches. BitMTP-2B reaches 964.2 token/s peak decode throughput, a $2.64\times$ speedup over BitNet b1.58-2B-SFT at comparable prediction accuracy. Code and data will be released.

We present BitDance, a scalable autoregressive (AR) image generator that predicts binary visual tokens instead of codebook indices. With high-entropy binary latents, BitDance lets each token represent up to $\mathbf{2^{256}}$ states, yielding a compact yet highly expressive discrete representation. Sampling from such a huge token space is difficult with standard classification. To resolve this, BitDance uses a binary diffusion head: instead of predicting an index with softmax, it employs continuous-space diffusion to generate the binary tokens. Furthermore, we propose next-patch diffusion, a new decoding method that predicts multiple tokens in parallel with high accuracy, greatly speeding up inference. On ImageNet 256$\times$256, BitDance achieves an FID of **1.24**, the best among AR models. With next-patch diffusion, BitDance beats state-of-the-art parallel AR models that use 1.4B parameters, while using **5.4$\times$** fewer parameters (260M) and achieving **8.7$\times$** speedup. For text-to-image generation and image editing, BitDance trains on large-scale multimodal tokens and generates high-resolution, photorealistic images efficiently, showing strong performance and favorable scaling. When generating 1024$\times$1024 images, BitDance achieves a speedup of over **30$\times$** compared to prior AR models. We release code and models to facilitate further research on AR foundation models.

Vision-language models (VLMs), such as CLIP, have shown remarkable transferability in the biomedical domain, and prompt tuning enables efficient adaptation with limited supervision. In practice, biomedical images exhibit pronounced acquisition-induced style shifts, especially across imaging sites and protocols, while cross-modality transfer further introduces compound shifts from different physical acquisition processes. However, existing prompt-tuning methods ignore acquisition-induced style shifts, making the prompts sensitive to source-specific style and limiting the transferability to unseen medical domains. To address this, we formulate acquisition-induced style shifts as mixture shifts, where each domain is viewed as a mixture of latent style components. Under this formulation, the target risk depends on both mixture weights and latent component risks, motivating prompt adaptation that accounts for style-specific variation rather than relying on a single source-specific prompt. In this paper, we propose SaMoE, a Style-aware Mixture-of-Experts Multimodal Prompt Learning framework for adapting VLMs to biomedical domains under acquisition-induced style shifts. To account for latent component risks in prompt adaptation, SaMoE represents multimodal prompts with latent style-conditioned prompt transformations and composes them through image-conditioned routing, enabling each image to induce style-adapted visual prompts. By extracting style cues from intermediate visual representations, the sparse expert composition adapts style-related prompt weights without domain labels and is injected across multiple VLM layers to align with hierarchical visual-language representations. Extensive experiments on 15 medical datasets across 12 modalities and 10 organs demonstrate significant improvements in both accuracy and generalizability over state-of-the-art methods.

Learning meaningful representations from cellular microscopy images remains a central challenge in computational biology and computer vision. Recent advances in self-supervised learning (SSL) have shown that scaling model and dataset size can improve the recovery of biological relationships such as gene–gene or compound–target associations. However, this progress has been driven primarily by scale rather than structure—requiring vast data and compute while offering no guarantee that the resulting embeddings capture underlying biology. We introduce NOVA, a structure-first framework for learning biologically aligned representations of cellular perturbations. NOVA integrates curated biological graphs directly into the training objective of a vector-quantised transformer, aligning embeddings of related genes and pathways through graph-based regularisation. By incorporating these relational priors during training, NOVA learns biologically faithful representations with orders of magnitude less data and compute than current state-of-the-art SSL models, while achieving superior performance on gene-gene retrieval benchmarks. This approach reframes representation learning for microscopy from a problem of scale to one of structure, offering a more systematic and efficient route toward biologically grounded embeddings.

Large language model (LLM) agents increasingly automate machine-learning model building, yet biomedical benchmarks emphasize question answering, reasoning, and tool use or cover narrow slices of biomedical ML coding. We introduce BioXArena, a biomedical machine learning (BioML) coding benchmark for task-specific model building on heterogeneous, often multi-modal datasets. It contains 76 end-to-end tasks across 9 domains: sequence, single-cell, structure, network biology, chemical biology, perturbation dynamics, phenotype--disease, imaging, and text-integrated tasks. Each task packages primary-source data with hidden labels, held-out graders, and biology-aware metrics on a unified 0-to-100 scale; agents write runnable code, train models, and predict private test samples. BioXArena emphasizes realistic data interfaces: most tasks require agents to integrate multiple input sources, including 46 tasks that combine distinct modality families. We compare 18 harness--backbone configurations: general-purpose coding, biomedical, and ML agents plus our Minimal runner, all under the same controlled resource configuration. Codex with GPT-6 Astra achieves the highest task-level average score of 73.91 among 18 configurations, closely followed by Claude Code with Opus 5.5 (73.79). Matched-backbone comparisons show higher observed scores for mature coding systems than for Minimal; seven Minimal backends characterize model differences under a fixed runner. Beyond the main leaderboard, ablations, robustness checks, and scaling, runtime, and failure-mode analyses characterize how backbones, scaffolds, budgets, and domains affect BioML coding performance, exposing modeling and execution gaps across agent and LLM backbone families.

Biomicroscopy imaging involves the complex coordination of multiple procedures, such as focusing, illumination, and staining. Traditional unitary methods struggle with dynamical environments, while hybrid models pre-trained on natural images tend to introduce artifacts for higher perceptual metrics. Notably, biological imaging is not a one-off task. It is an iterative process similar to human expert’s observation, verification and adjustment. They can benefit from history, remembering which processing strategies are suitable and can also identify issues such as artifacts and mis-staining from verification. To this end, we propose BioMicroAgents, the first multi-agent framework with co-evolutionary capabilities designed to simulate and optimize whole biomicro imaging process including brightening, denoising, deblur, super-resolution, and virtual staining. Through reinforcement learning, the model can reflect based on historical experience to reduce artifacts and mis-staining, while enhancing edge and structural features to improve the accuracy of downstream tasks. Specifically, BioMicroAgents includes: (i) Multi-modal Analytical Agents, which utilize multimodal RAG to provide template images and biomicro knowledge, and employ the memory mechanism to analyze optimal processing paths; (ii) We propose three biomicro fidelity metrics and utilize three agent groups (MicroNavigator, MicroProcessor, and MicroAuthenticator) to perform rigorous analysis, processing, and verification. (iii) Co-evolutionary ability, which achieves dynamic policy and self-correction through interactive feedback strategies and memory mechanisms, allowing the agent to evolve from history and other agents. We test model across 12 biomicro tasks, BioMicroAgents achieved improvements in perceptual quality while adhering to strict fidelity in sub-cellular features. It also enhances the accuracy of downstream tasks such as segmentation and cell counting, showcasing the potential of general-purpose visual agents in bioscience.