NeurIPS 2026 Papers

Humans know when to reach for help e.g. $347 \times 28$ warrants a calculator while $2 + 2$ does not. Language models, by default, do not. Prompt-based approaches can instruct a model when to invoke tools, but this external scaffolding does not teach the model to recognize the boundary of its own knowledge. Reinforcement learning approaches that assign a single outcome reward to the whole trajectory fare no better: trajectory-level credit cannot isolate which tool call in a successful episode actually helped, nor penalize unnecessary calls. We propose \textbf{CARL} (\textbf{C}ompetence-\textbf{A}ware \textbf{R}einforcement \textbf{L}earning), which trains a critic on the model's own rollouts to learn where the model's parametric knowledge suffices and where it needs external help. By decomposing each rollout at natural tool-use boundaries (e.g., code fence delimiters and context block transitions), CARL assigns independent credit to each segment from a single binary outcome, without external judges or step-level annotations, addressing the credit-assignment limitations of trajectory-level methods. As a result, erroneous tool calls, incorrect extractions, and unnecessary calls each receive appropriately signed advantages under a single scheme. We show quantitatively and qualitatively that the trained critic captures the model's domain competence: it separates parametrically solvable from tool-dependent questions with AUC 0.93 at 7B. On five benchmarks spanning arithmetic, multi-hop factual QA, and numerical reasoning over financial tables, CARL improves exact-match accuracy by 6.7 points at 7B and 10.6 points at 3B over the strongest trajectory-level baseline (Search-R1 PPO), with the largest gain (+8.3 EM at 7B, +9.0 EM at 3B) on Musique, the most compositional multi-hop benchmark. Compared to a representative trajectory-level baseline (Search-R1 GRPO), the model issues 56\% fewer tool calls on questions answerable from parametric knowledge while remaining ${\sim}10$ EM points more accurate on those same questions, an emergent consequence of the critic learning where this model's competence ends. Gains are largest at small scale: the 3B model improvement is $1.6\times$ the 7B improvement, suggesting that knowing when to ask for help disproportionately benefits models with smaller parametric memory.

At deployment, a multivariate forecaster may expose only a recent input window and a frozen point prediction, with no access to ground truth, calibration residuals, repeated inference, or model internals. We study deployment-constrained forecast risk localization: ranking future timesteps by likely error before they are observed. We introduce SubDx, a zero-parameter diagnostic that audits the forecast against the local cross-channel subspace of the input history. Each predicted channel vector receives an off-subspace residual score, computed with no labels, no learned parameters, and less than 1% runtime overhead on Traffic. Under a linear factor model, we derive closed-form score distributions and a two-population AUROC expression showing how heterogeneous off-subspace error visibility makes the ranking identifiable. On Traffic (N=862), SubDx reaches within-horizon AUROC 0.90 and reduces MSE by 62% at 50% coverage; on Solar (N=137), it reaches AUROC 0.85 with 67% MSE reduction. The same post-hoc score applies to frozen pretrained forecasters, with Chronos reaching AUROC 0.838 on Traffic without fine-tuning. Across trained backbones, pretrained models, and synthetic controls, the results show that local cross-channel consistency can provide a practical abstention signal for redundant multivariate systems.

Long-term conversational memory poses a fundamental system challenge for LLM agents: reasoning over all past interactions improves coverage but incurs prohibitive token costs and noise, while retrieval-based methods are efficient yet often fail on multi-hop and temporal questions. Existing multi-step retrievers partially address this gap, but typically operate over an ever-growing textual history, causing context expansion and noise accumulation across iterations. We propose MemR$^3$, a backend-agnostic closed-loop controller that reformulates conversational memory retrieval as explicit stateful decision-making. At each iteration, MemR$^3$ maintains an evidence-gap state consisting of grounded information (evidence) and unresolved information requirements (gaps), and uses this state to route among three actions: retrieve, reflect, and answer. Newly retrieved snippets are incorporated into the state and then masked from subsequent prompts, so the model conditions on a compact summary of progress rather than the full retrieval history. This design turns multi-step memory retrieval into a bounded, inspectable control process whose query-time context scales with the state instead of raw accumulated text. Experiments on LongMemEval$_s$ show that MemR$^3$ surpasses implicit multi-step retrieval baselines and often approaches or exceeds full-context prompting while using only 1%-5% of the full-context tokens on long conversations. On LoCoMo, MemR$^3$ consistently improves over its underlying RAG- and Zep-based memory backbones. These results suggest that explicit evidence-gap tracking is an effective abstraction for token-efficient conversational memory.

Large language models (LLMs) have shown promise in scientific discovery tasks such as protein optimization, but remain limited by the lack of a mechanism for persistent and reusable knowledge accumulation. Existing approaches adapt LLMs through parameter optimization or trajectory refinement, but in both cases, knowledge remains implicit, either encoded in opaque model weights or scattered across transient solution histories. As a result, knowledge cannot be systematically inspected, attributed, or reused. We propose \emph{KnowEvo}, a framework that treats externalized knowledge itself as the optimization target. Our key insight is that scientific knowledge is inherently modular: multiple interacting mechanisms jointly determine the design outcomes. This makes knowledge evolution fundamentally challenging, as improvements cannot be attributed to individual components in isolation. To address this, \emph{KnowEvo} represents knowledge as executable expertise blocks and organizes their refinements into strategy mutation trees, enabling structured exploration over a combinatorial knowledge space. The framework combines a forward exploration phase that adaptively allocates evaluation budget across strategy variants with a backward revision phase in which an agentic sandbox converts rollout evidence into localized program improvements. This design enables persistent accumulation, attribution, and refinement of mechanistic expertise across iterative experimentation. We evaluate \emph{KnowEvo} on protein optimization tasks including thermostability and solubility. Results show that knowledge evolution consistently improves over initial expert heuristics and outperforms monolithic program evolution and prior LLM-guided evolution baselines across ranking metrics such as pass@$k$, MRR, MAP, and NDCG@5. Codes are available at \url{https://anonymous.4open.science/r/KnowEvo-6E3B/}.

Cooperative multi-agent reinforcement learning under task variability requires both identifying which task an episode belongs to and adapting the learning algorithm to that identity. Recent class-aware methods address the first half of this problem by clustering trajectories into latent task classes and feeding the class label to the agent policy. They leave the second half largely untouched: the Bellman update and the credit assignment performed by the value mixer remain identical across all task classes. We show empirically that, even when task identity is recovered with high accuracy, the Bellman dynamics across classes diverge sharply, with simple classes saturating early and harder classes sustaining large temporal-difference residuals and large cross-agent value disagreement. Building on this observation, we propose Know Your Task, Learn It Right (KYT), a task-aware extension of value factorization that injects per-class learning signal into the Bellman update itself. KYT maintains a lightweight tracker that summarizes the per-class TD-error and cross-agent value spread without any extra network, and uses the tracker in two places: an optimistic target that adds a bounded per-class exploration bonus, and a class-conditioned mixer that modulates credit assignment through a softplus-gated affine layer preserving the IGM property. Across StarCraft II micromanagement benchmarks including SMACv2 and the unit-combination suite SurComb, KYT outperforms strong class-aware and exploration-aware baselines, with the largest gains concentrated on the hardest task classes that current methods underfit.

Memory enables long-context dialogue agents to maintain user state across many turns, but it also introduces a fundamental source-choice problem: grounding the response in the immediate context or retrieved prior memory? A wrong choice incurs two complementary costs: neglecting valid aged-out evidence, or overwriting sufficient context with stale memory. Existing systems and benchmarks largely conflate retrieval with source choice, relying on the model to resolve it implicitly and leaving this critical decision largely unexplored. In this paper, we introduce MSBench, a controlled benchmark that isolates source choice by pairing questions answerable exclusively from context with those requiring prior memory. Under this protocol, strategies that depend on the model's implicit judgment exhibit sharply degraded performance, exposing the limits of leaving source choice ungoverned. To this end, we propose the Memory-Source Reasoner (MSR), which frames source choice as an explicit metacognitive decision: it reasons about contextual sufficiency, answers directly when appropriate, and selectively retrieves missing evidence otherwise. Experimental results show that MSR achieves the best overall accuracy and selection accuracy across four answer backbones. On the primary GPT-5-mini run, it improves over the best non-MSR baselines by 12.6 points in overall accuracy and 6.3 points in selection accuracy. Our benchmark, code, prompts, and evaluation scripts are available at https://anonymous.4open.science/r/msbench-msr-0C72.

Post-LayerNorm Transformers have seen limited practical adoption due to enduring difficulties in scaling them to large depths. While prior research has focused on stabilising Post-LayerNorm by improving conditioning at initialisation, stability often deteriorates when models are trained with large learning rates, forcing additional compromises. In the decoder-only settings we study, existing Post-LayerNorm methods fail to outperform strong Pre-LayerNorm baselines. We propose KiteNorm, a novel normalisation method designed to break this trend. KiteNorm achieves stability through residual scaling and a regularisation technique based on hidden-state variance. Beyond stabilisation, KiteNorm learns separate scales for the skip and residual branches in each sublayer, improving performance. Across decoder-only Transformers up to 1B parameters, KiteNorm remains stable throughout training and consistently outperforms leading baselines, with scaling laws favouring KiteNorm across depth, width, batch size, and training budget. Ablations show that each component is necessary for the full stability and performance gains. We offer theoretical insights into KiteNorm through a new perspective on the gradient vanishing problem in Post-LayerNorm, linking the underlying hidden-state expansion to rank collapse.

Metric-induced discrete flow matching (MI-DFM) exploits token-latent geometry for discrete generation, but its practical use is limited by two issues: heuristic schedulers requiring hyperparameter search, and finite-step path-tracking error from its first-order continuous-time Markov chain (CTMC) solver. We address both issues. First, we derive a kinetic-optimal scheduler for prescribed scalar-parameterized probability paths, and instantiate it for MI-DFM as a training-free numerical schedule that traverses the path at constant Fisher–Rao speed. Second, we introduce a finite-step moment correction that adjusts the jump probability while preserving the CTMC jump destination distribution. We validate the resulting method, GibbsTTS, on codec-based zero-shot text-to-speech (TTS). Under controlled comparisons with a unified architecture and large-scale dataset, GibbsTTS achieves the best objective naturalness and is preferred in subjective evaluations over masked discrete generative baselines. Additionally, in comparison with the evaluated state-of-the-art TTS systems, GibbsTTS shows strong speaker similarity, achieving the highest similarity on three of four test sets and ranking second on the fourth. Demo and code are available in the supplementary materials.

Modern Bayesian optimization and adaptive sampling methods increasingly rely on nonlinear parametric models, yet theoretical guarantees for such models under adaptive data collection remain limited. Existing analyses largely focus on Gaussian processes, kernel machines, linear models, or linearized neural approximations, leaving a gap between theory and the nonlinear models used in practice. We develop a kernel-based framework for analyzing regularized nonlinear parametric models trained on adaptively collected data. Our approach uses kernels over the parameter space to induce reproducing-kernel Hilbert space structures over the corresponding model class, yielding confidence bounds for models trained with broad classes of regularized convex losses. We show how these bounds can support convergence guarantees for nonlinear acquisition and surrogate models, including randomized regularized policies that select points by maximizing a trained random model. These results provide a unified route to analyzing nonlinear parametric models in Bayesian optimization and related adaptive optimization settings.

We propose kernel-gradient drifting, a one-step generative modeling framework that replaces the fixed Euclidean displacement direction in drifting models with directions induced by the kernel itself. Standard drifting is attractive because it enables fast, high-quality generation without distilling a large pretrained diffusion model, but its theory is currently understood mainly for Gaussian kernels, where the drift coincides with smoothed score matching and is identifiable. Our gradient-based reformulation exposes this score-based structure for general kernels: the resulting drift is the score difference between kernel-smoothed data and model distributions, yielding identifiability for characteristic kernels and a smoothed-KL descent interpretation of the drifting dynamics. Since kernel gradients are intrinsic tangent vectors, the same construction extends naturally to Riemannian manifolds and to discrete data via the Fisher-Rao geometry of the probability simplex. Across spherical geospatial data, promoter DNA and molecule generation, kernel-gradient drifting enables state-of-the-art one-step generation beyond the Euclidean setting without distillation.

Value overestimation is a central challenge in offline reinforcement learning (RL), arising when Bellman backups propagate out-of-distribution (OOD) errors that lead to catastrophic inflation of $Q$-values. Prior approaches address this issue by encouraging the learned policy to stay close to the behavior policy, but they often result in suboptimal policies or require significant training complexity. More recent work attempts to tackle this issue through minimal modifications to the critic network, yet it remains limited in either theoretical grounding or practical implementation. In this work, we present Kernel Value Regression (KVR), a simple critic regularization method based on kernel regression. KVR builds on the vanishing extrapolation property of kernel ridge regression, which encourages value estimates to decay toward zero away from the data support. This inductive bias naturally suppresses $Q$-values for OOD actions without auxiliary regularizers, providing a structural mechanism for addressing value overestimation. Using random Fourier features, KVR can be incorporated into standard actor-critic algorithms with only minimal modifications. We empirically demonstrate that these properties persist in complex offline RL environments, leading to strong performance across diverse tasks in the D4RL and OGBench benchmarks.

The Maximum Mean Discrepancy (MMD) is a cornerstone statistic for nonparametric two-sample testing, but its test power is dictated entirely by the chosen kernel. Because any fixed kernel inherently fails to distinguish certain distributions, the kernel must be dynamically optimised. However, data-driven optimisation violates the foundational i.i.d. assumption, forcing a strict trade-off in existing frameworks. Ratio criteria ignore this dependence, inducing overfitting and variance collapse on rich kernel classes. Conversely, aggregation methods bypass the dependence using finite grids, but this strategy cannot scale to continuous search spaces like deep kernels. To break this dichotomy, we establish data-driven kernel selection as a model selection problem. We propose Complexity-Penalised MMD (CP-MMD), a criterion derived from a uniform concentration inequality (extending Maurer's framework to the two-sample setting) that bounds the post-optimisation distribution by penalising the empirical MMD with the complexity of the kernel search space. Because this penalty mathematically absorbs the cost of optimisation, CP-MMD enables direct, grid-free maximisation over continuous parametric classes, including scalar bandwidths, polynomial coefficients, and deep network parameters. By formally accounting for optimisation complexity, we theoretically guarantee that CP-MMD maximises true test power while ensuring unconditional Type-I validity. Consequently, CP-MMD enables grid-free kernel selection across linear, polynomial, and deep regimes, matching or exceeding state-of-the-art test power.

Operator learning has experienced a shift toward attention-based architectures. This trend has more recently extended to physics-informed operator learning, where stacked cross-attention blocks add considerable computational cost relative to the foundational PI-DeepONet. We ask whether this complexity is necessary. Starting from the DeepONet branch-trunk decomposition, we replace the bilinear readout with a single softmax-weighted readout in which the trunk emits a query vector and the branch emits a set of key-value pairs—a structure we show admits universal approximation via reduction to a simplex-bilinear form. The resulting architecture, KerONet, uses no iterated attention, layer normalization, or explicit learnable projection matrices. Under a unified hyperparameter optimization protocol with a strict $200{,}000$-parameter budget, we evaluate KerONet alongside iterated-attention architectures (PIT, PINTO) and PI-DeepONet on four canonical 1D PDEs whose solutions range from globally smooth to sharply localized. KerONet outperforms PIT and PINTO on every PDE, across all five seeds, in accuracy (approximately $2$–$4\times$ lower rel-$L^2$ error), reproducibility ($2$–$31\times$ lower standard deviation), and speed ($1.8$–$4.4\times$ faster training time). Against PI-DeepONet, KerONet shows substantial gains where input-dependent localized features dominate ($2.5$–$8\times$ lower error on Burgers and Allen-Cahn) while exhibiting comparable accuracy on smoother operators (advection and diffusion-reaction)–a balance the iterated-attention architectures do not achieve. Our findings demonstrate that iterated cross-attention is not necessary to realize the benefits of query-dependent feature mixing in physics-informed operator learning.

Training-free safeguards for text-to-image diffusion models often rely on a reusable safety signal, such as an unsafe direction or global toxic subspace, applied broadly across prompts. We provide a controlled geometric analysis of this global-unsafety assumption and reveal a consistent coverage-selectivity trade-off: compact unsafe subspaces fail to cover heterogeneous unsafe semantics, whereas broader aggregation increasingly distorts safety-adjacent benign prompts. Motivated by this finding, we propose CALM (Counterfactual Adaptive Local Modulation), a training-free safeguard that replaces uniform global removal with prompt-local counterfactual correction. Using matched unsafe-benign anchors, CALM routes each prompt to active unsafe categories, minimally edits only violating token representations toward the safe side, and suppresses positively aligned unsafe residual components. Across broad evaluation, CALM improves unsafe-content suppression while preserving benign utility, demonstrating that local counterfactual correction provides a more selective alternative to global unsafe-signal removal.

LLMs are increasingly applied to cybersecurity workflows, where they are expected to translate analysts' intent into tool invocations. However, existing evaluations focus on knowledge-based assessments or end-to-end agentic tasks, and do not directly measure LLMs’ ability to generate executable commands for real-world cybersecurity tools. This gap is critical because cybersecurity operations rely on strict command-line interfaces (CLIs), where minor syntax errors, incorrect flag-value bindings, or argument misordering can invalidate execution. We introduce KaliBench, a fine-grained benchmark and dataset for natural-language-to-CLI translation on Kali Linux, comprising 8,504 query-command pairs spanning 1,642 tools across 23 capability dimensions and 5 security phases. KaliBench is constructed via a manuscript-grounded pipeline with deterministic canonicalization and alias-aware evaluation, enabling precise and reproducible assessment of tool selection and argument construction. To ensure both semantic correctness and practical executability, we develop a multi-stage verification pipeline that combines LLM-based validation, sandboxed terminal execution, and human-in-the-loop refinement. Building on these fine-grained, deterministic signals, KaliBench further enables runtime-free verifiable rewards for training. Across three evaluation modes and over 20 general-purpose and security-focused LLMs, no model exceeds 35\% exact-command accuracy in the unrestricted setting, highlighting the difficulty of accurate CLI-based cybersecurity tool use without explicit tool hints. We further show that supervised fine-tuning and reinforcement learning with verifiable rewards derived from KaliBench significantly improve an 8B model and achieve performance comparable to a 685B MoE model. KaliBench will be publicly released.

Inherent temporal heterogeneity, such as varying sampling densities and periodic structures, has posed substantial challenges in zero-shot generalization for Time Series Foundation Models (TSFMs). Existing TSFMs predominantly rely on massive parameterization to absorb such heterogeneity, as their static tokenization and positional encoding schemes entangle diverse temporal patterns into a fixed representation space, encouraging memorization rather than adaptation. To address this limitation, we propose Kairos, a flexible and parameter-efficient TSFM dedicated to forecasting tasks, which decouples temporal heterogeneity from model capacity through a novel tokenization perspective. Kairos introduces a dynamic patching tokenizer and a mixture-of-size encoding that adapt observational granularity to local information density, enabling fine-grained temporal abstraction without increasing model width or depth. In addition, we design a multi-granularity positional embedding based on dynamic rotary encodings, which conditions on instance-level spectral features and temporal structure induced by dynamic patching tokenization, allowing robust modeling of diverse temporal dependencies. Trained on a novel Predictability-Stratified Time-Series (PreSTS) corpus, Kairos achieves superior zero-shot performance with substantially fewer parameters on two mainstream benchmarks, GIFT-Eval and Time-Series-Library.

In retrieval-augmented generation (RAG) scenarios, the prefill delay of the large language model (LLM) long inputs is the dominant component of time-to-first-token (TTFT). Reusing KV caches across different LLM inputs reduces this cost, but the resulting cross-context mismatch on reused segments has a significant degradation on answer quality via propagating through prefill attention into decode-time logits. Existing selective-recomputation methods mitigate this issue by recomputing tokens chosen by empirical importance heuristics, which overlook the propagation of cache mismatches into answer-side errors, thereby failing to effectively strike the quality-TTFT tradeoff. To address this issue, we propose a perturbation-theoretic KV cache reuse framework, KVFocus, which selects reused KV caches through a perturbation-theoretic token-risk score during the prefill process. Specifically, we first develop a first-order analysis that traces reuse error in three stages: source mismatch on the reused segment, suffix contamination during prefill, and decode-time logit stability. This analysis reveals that the leading-order answer-side error decomposes into a product of source-side V-drift and downstream suffix-to-segment attention concentration. Guided by this multiplicative structure, we define a token-risk score that operationalizes the bound during the inference. A plug-and-play selector for KV cache reuse in existing RAG-based LLM serving is then designed by recomputing the top-$r\%$ tokens directly suppresses the predicted leading-order answer-side error, covering both sides of the pathway that prior one-sided heuristics miss. Empirical results across three 7-8B instruction-tuned LLMs (Mistral, Qwen2.5, Qwen3) and four multi-hop QA benchmarks demonstrate that the proposed KVFocus achieves a superior trade-off between the ansewer quality and the TTFT in comparison to state-of-the-art baselines, with up to $3\times$ TTFT speedup over full prefill at matched answer quality.

rocess Reward Models (PRMs) have been proven to be highly effective in guiding test-time scaling (TTS) methods, which significantly boost the capabilities of LLM-based multi-agent systems. However, existing PRMs are *text-based*: they re-encode the entire trajectory text from scratch. In long multi-agent rollouts, the scoring cost, growing quadratically with respect to sequence length $L$, creates a severe computational bottleneck, severely limiting PRMs' application in long-context scenarios. To resolve this, we introduce KV-PRM, a highly efficient process reward model that eliminates the heavy text re-encoding by directly reading the KV cache produced naturally during the LLM's generation phase. By processing a single "verify token" against the pre-existing KV cache, KV-PRM reduces the scoring cost from $O(L^2)$ to $O(L)$. We formally prove that the KV cache contains strictly greater information capacity than text, and is more efficient for downstream reward modeling. Empirically, across the MATH, GSM8K, and AIME benchmarks, KV-PRM matches or strictly outperforms text-PRMs under various TTS methods such as Beam Search, MCTS, and Weighted Voting, with up to a *5,000*$\times$ reduction in scoring FLOPs, a *37*$\times$ reduction in latency, and a *34*$\times$ reduction in per-sequence memory footprint compared to text-based PRMs.

Large Language Models (LLMs) rely heavily on Key-Value (KV) caching to minimize inference latency. However, standard KV caches are context-dependent: reusing a cached document in a new context requires recomputing KV states to account for shifts in attention distribution. Existing solutions such as CacheBlend, EPIC, and SAM-KV mitigate this issue by selectively recomputing a subset of tokens; however, they still incur non-negligible computational overhead (FLOPs) and increased Time-to-First-Token (TTFT) latency. In this paper, we propose KV Packet, a recomputation-free cache reuse framework that treats cached documents as immutable ``packets'' wrapped in light-weight trainable soft-token adapters, which are trained via self-supervised distillation to bridge context discontinuities. Experiments on Llama-3.1 and Qwen2.5 demonstrate that the proposed KV Packet method achieves near-zero FLOPs and lower TTFT than recomputation-based baselines, while retaining F1 scores comparable to those of the full recomputation baseline. Code and experimental results are available at https://anonymous.4open.science/r/kv_packet_anonymous-46FC/.