NeurIPS 2026 Papers

Low-Rank Adaptation (LoRA) has emerged as a widely used technique for adapting large language models (LLMs) to new domains, due to its modular design and broad availability on platforms such as HuggingFace. This availability has motivated efforts to reuse existing LoRAs for domain generalization. However, existing methods often rely on explicit task labels or additional training, which are impractical for deployment. Moreover, they typically activate a fixed number of entire LoRA, leading to parameter redundancy or insufficiency that degrade performance. In this paper, we propose HiLoRA, a training-free framework that performs adaptive hierarchical routing over LoRA pools. Drawing on structural properties of LoRA, we define rank-one components (ROCs), where each LoRA rank is treated as an independent unit. For a given input, HiLoRA first adaptively selects a subset of LoRAs and determines their ROC allocation based on Gaussian likelihoods at the sequence level. At the token level, it refines routing by activating only the most informative ROCs. We further provide theoretical guarantees that HiLoRA selects the most relevant LoRAs with high probability. Extensive experiments show that HiLoRA achieves substantial improvements in domain generalization, with accuracy gains of up to 70\% over state-of-the-art baselines, while maintaining comparable inference throughput.

Mechanistic interpretability aims to explain a model’s behavior by identifying causally responsible internal structures. Dictionary-based explainers such as sparse autoencoders and transcoders are a primary tool, but their faithfulness under out-of-distribution (OOD) shift has received little systematic attention. We show that distribution shift rotates the subspace that the model actively uses, misaligning the explainer’s dictionary trained on in-distribution (ID) activations. We formalize this misalignment as the faithfulness gap, a geometric distance between the ID dictionary and the OOD-active subspace, and show that it controls OOD faithfulness degradation. To reduce this gap, we propose the Geometry-Adaptive Explainer (GAE), which realigns the explainer's dictionary with the OOD-active subspace while preserving the original feature structure. This requires only unlabeled OOD activations and no gradient updates. We prove that GAE improves over the unadapted ID explainer, with excess loss bounded quadratically by the second-moment shift. Empirically, GAE even matches or surpasses all training-based baselines in causal faithfulness across multiple models and OOD settings.

Not all mistakes are alike. For example, in medical diagnosis, false positive and false negative mistakes have a different impact and should not be aggregated. By revisiting the classical theory of boosting through this lens, we introduce the notion of $Z$-learner, which returns weak hypotheses achieving joint false-positive/false-negative guarantees $(FP,FN)$ from some set $Z \subseteq [0,1]^2$. This formulation generalizes the standard notion of $\gamma$-weak learning — recovered by taking $Z=$ { $(z_1,z_2): z_1+z_2<1/2-\gamma$ } — and allows for more flexible notions of weak learnability that account for the impact of the distribution on the achievable learning guarantees. This raises a natural question: which sets Z are boostable? We answer this by giving a complete geometric characterization that unifies, as special cases, both the classical boosting dichotomy and the recent multi-objective extension of Bressan et al. (COLT 2025). We also present an adaptive algorithm for boosting $Z$-learners without relying on the prior knowledge of~$Z$, generalizing AdaBoost's ability to handle unknown margin parameters. Finally, we demonstrate that $Z$-learners can achieve confidence amplification. However, the analysis is significantly more involved than in classical boosting and relies on the centerpoint theorem from discrete and computational geometry.

Selective state space models (SSMs) such as Mamba-3 are efficient and increasingly competitive with Transformers. However, their fixed recurrent state remains a bottleneck for using context as temporary memory. To address this issue, we propose $ \textit{GEAR} $ (GEnerator-Adaptive state space models for associative Recall), a slow-hypernetwork that produces coefficients for LoRA updates to the SSM token-to-recurrence parameter generators. GEAR makes the parameter generator context-adaptive while preserving the original linear-time Mamba-3 scan. To evaluate whether this adaptation improves in-context recall, we pretrain 180M-parameter models for 2B FineWeb-Edu tokens and compare against Mamba-3 baseline and ablations at the same scale. In the Multi-Query Associative Recall (MQAR) task, models see $N$ in-context key-value bindings and must predict the value paired with a queried key. In compact MQAR evaluation where the adaptive path is active at the query, GEAR improves candidate-restricted NLL over Mamba-3 by 0.70, 0.69, and 0.48 at $ N=32,64,128 $, respectively. Beyond compact MQAR, GEAR reduces the penalty on older queried bindings and gives the clearest gain within the Mamba-3 family under moderate random-filler interference. These results show that context-adaptive SSM generator improves associative recall while preserving Mamba-3 language modeling capabilities.

Functional optimization problems are typically solved by optimizing the parameters of a fixed representation, such as a neural network, resulting in highly nonconvex losses that complicate both training and theoretical analysis. An interesting alternative is functional gradient descent (FGD) --- that is, gradient descent directly in function space --- which benefits from strong convergence results and admits a clean theory. However, FGD is difficult to implement in practice because functional gradients are infinite-dimensional, and thus cannot be fully computed nor stored in memory. Existing implementations therefore rely on fixed approximations, which introduce approximation error. We propose a new, theoretically-grounded FGD algorithm that adapts the representation of the functional gradients over the course of optimization. By explicitly incorporating this approximation into the analysis, we establish convergence to a stationary point (for smooth losses) and to a global minimizer (under smoothness + a Polyak-Łojasiewicz-type condition) regardless of our approximations. To the best of our knowledge, this is the first implementable FGD method with such guarantees in a general setting. We demonstrate the effectiveness of our method on regression, numerical solution of PDEs, and modern computer vision. Across settings, our method consistently outperforms both FGD with fixed approximations and neural network baselines in efficiency and accuracy.

Modern optimizers, like Muon, impose matrix-wise geometry constraints on their updates. These matrix-wise constraints can be unified under Linear Minimization Oracle (LMO) theory. However, all current methods impose fixed LMO geometries for the update rules, chosen by-design or empirically, which are not necessarily optimal according to the problem's geometry. We introduce a novel efficient data-driven criterion for dynamically choosing proxy-optimal update LMO geometries on individual Deep Neural Network layers. Derived in closed form from gradient and activation statistics using a single-step random feature regression surrogate model, our criterion navigates a design space interpolating from SGD to Muon updates. Moreover, integrating parameter-wise preconditioning allows our framework to recover SGD, Muon, Adam, and MuAdam as specific extrema. To make this adaptive approach scalable, we pair it with efficient computational strategies, achieving only a $\sim 3$% runtime overhead on highly optimized baselines. As a proof of concept, we show that this data-driven optimizer matches or exceeds the performance of the best performing optimizer between Muon and AdamW across three different training scenarios. Ultimately, this work provides evidence that LMO geometry can be successfully and efficiently adapted from runtime data, opening a new pathway for optimizer design beyond static geometries.

Agentic forecasting is important for decision-making in dynamic environments, but it remains challenging because agents must reason from incomplete, time-limited evidence and produce calibrated probabilities before outcomes are resolved. Memory provides a natural mechanism for transferring experience from resolved forecasts to future prediction tasks. However, existing agent-memory methods are not tailored to forecasting, as they typically store past interactions, reflections, or factual associations without explicitly representing reusable predictive factors or calibration knowledge. We propose ForecastCompass(FoCo), an adaptive factor-based memory framework for agentic forecasting. FoCo organizes forecasting experience with a hierarchical forecasting-task taxonomy, enabling retrieval task-relevant forecasting knowledge. It maintains two complementary memory components: factor memory, which captures reusable predictive dimensions, and reasoning memory, which encodes probability updating, uncertainty handling, and calibration principles. Using retrospective analyses as learning signals, FoCo iteratively revises memory through a verbalized memory-revision procedure, enabling the agent to accumulate transferable forecasting knowledge over time. Experiments on Prophet Arena and FutureX with GPT-5-mini and Gemini-2.5-Flash show that FoCo improves both probabilistic accuracy and calibration.

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.

Federated learning (FL) encounters substantial challenges due to heterogeneity, leading to gradient noise, client drift, and partial client participation errors, the last of which is the most pervasive but remains insufficiently addressed in current literature. In this paper, we propose FedAdaVR, a novel FL algorithm aimed at solving heterogeneity issues caused by sporadic client participation by incorporating an adaptive optimiser with a variance reduction technique. This method takes advantage of the most recent stored updates from clients, even when they are absent from the current training round, thereby emulating their presence. Furthermore, we propose FedAdaVR-Quant, which stores client updates in quantised form, significantly reducing the memory requirements (by 50%, 75%, and 87.5%) of FedAdaVR while maintaining highly competitive model performance. We analyse the convergence behaviour of FedAdaVR under general nonconvex conditions and prove that our proposed algorithm can eliminate partial client participation error. Extensive experiments conducted on multiple datasets, under both independent and identically distributed (IID) and non-IID settings, demonstrate that FedAdaVR consistently outperforms state-of-the-art baseline methods.

Vision-Language-Action (VLA) models have shown strong potential for general robotic manipulation, but contact-rich tasks still require timely action refinement using force/torque feedback. Existing force-aware VLA models usually align all modalities to a single low operating frequency, which discards high-frequency contact cues that are important for reactive action correction. Additionally, they use static fusion, which cannot adaptively balance visual context and force feedback across manipulation processes. To mitigate these issues, we propose FAVLA, a force-adaptive multi-rate VLA model that explicitly separates low-frequency visual-language reasoning from high-frequency force-conditioned action refinement. Based on the $\pi_0$-style VLM-action expert architecture, FAVLA uses a low-rate VLM to encode visual-language-force context and predict near-future force statistics, while a high-rate action expert refines action chunks using the latest force observations. To decide \emph{when} to update actions, we introduce a Force-Adaptive Multi-rate Inference (FAMI) mechanism, which schedules the action expert's inference rate from predicted future force variance, keeping low-rate reasoning in stable phases and increasing update frequency near contact transitions. To decide \emph{how} to use force, we design a Force-Guided Dynamic Fusion (FGDF) module, which injects high-frequency force features into the action expert and dynamically balances visual-semantic and force cues across manipulation stages. Extensive real-robot experiments on high-precision and contact-rich tasks show that FAVLA outperforms force-aware VLA baselines, achieving an average success rate of 88.8\% and lower peak contact forces during manipulation.

Structure-informed protein representation learning is essential for effective protein function annotation and $\textit{de novo}$ design. With the rapid progress of protein structure prediction and experimental structure determination, a central challenge in the field is shifting from merely obtaining structural models to reliably leveraging predicted or measured structures for downstream learning. However, both crystal and AlphaFold-like predicted structures can contain structural uncertainty, which can bias local neighborhood construction and make robust protein representation learning challenging. To address these issues, we propose a novel equivariant Transformer-State Space Model (SSM) hybrid framework, termed $E^3$former, designed for efficient and robust protein representation. Our approach uses energy function-based receptive fields to construct proximity graphs that adaptively stabilize neighborhood selection under structural deviations, and incorporates an equivariant high-tensor-elastic selective SSM within the transformer architecture. These components allow the model to adapt to complex geometric interactions and extract structural features with a higher signal-to-noise ratio. Empirical results demonstrate that our model outperforms existing methods in structure-intensive tasks, such as inverse folding and binding site prediction, particularly when using predicted structures, owing to its enhanced tolerance to data deviation and structural uncertainty. Our approach offers a novel perspective for conducting biological function research and drug discovery using imperfect but increasingly available protein structure data. Our code is available on \url{https://anonymous.4open.science/r/E3former-207E}.

Learning effective policies for adaptive data acquisition remains challenging: posterior-based methods rely on surrogate models and posterior approximations that can be misspecified or biased, while direct policy-learning methods map from historical observations and fail to exploit available model representations, making learning harder. We introduce policy learning with belief representations (POLAR), based on the insight that optimal data acquisition depends on the observation history only through a sufficient belief state. Specifically, POLAR decouples representation learning from policy learning by leveraging pretrained predictive foundation models as belief-state encoders, training a policy head on top of their representations. This yields a simple, unified amortised policy learning framework for Bayesian experimental design, Bayesian optimisation, and active learning, differing only in the task-specific utility used to train the policy. Empirically, we find that POLAR outperforms state-of-the-art amortised methods across diverse tasks while requiring far fewer training samples, demonstrating a significant step in the scalability and efficiency of amortised data acquisition.

Reinforcement Learning with Verifiable Rewards empowers Large Language Models to enhance their reasoning capabilities. Most algorithms rely on on-policy training, which is stable but inefficient due to its strict dependency on fresh samples. To overcome this inefficiency, off-policy RL decouples experience generation from policy optimization via Importance Sampling (IS). However, this introduces a policy discrepancy between the behavior policy and the target policy. We observe that mainstream RLVR algorithms struggle to maintain a balance between performance and stability in off-policy scenarios, suffering from either severe performance degradation or training collapse. We identify the root cause as the failure of static trust regions to maintain the bias--variance trade-off associated with IS ratios in off-policy settings. To address this, we propose **DI**stribution-adaptive **P**olicy **O**ptimization (**DIPO**), a novel hyperparameter-free RLVR algorithm employing an adaptive trust region. DIPO leverages the intrinsic statistics of the IS ratio distribution to achieve a robust bias--variance trade-off. Experimental results demonstrate that DIPO enables stable training in extreme off-policy scenarios, achieving performance comparable to or even surpassing the on-policy baseline, and outperforming other off-policy baselines by 13.9%.

Residual connections with PreNorm are the default design in modern transformers due to their strong optimization stability, but recent research has shown that they weaken the functional differentiation of deeper blocks and increase representational redundancy in later layers. This issue also appears in diffusion models using similar architectures, where the deep block is expected to support progressive denoising refinement rather than repeated computation over similar states. Motivated by this, we revisit cross-layer shortcuts in diffusion transformers and argue that their role should go beyond merely transferring shallow information, instead guiding deeper blocks to improve depth utilization. To this end, we propose \textit{\textbf{Res}idual Adaptive Layer \textbf{Norm}alization} (\textbf{ResNorm}), which routes residual cross-layer shortcuts into the existing adaLN modulation. ResNorm preserves the optimization advantages of PreNorm residual connections, while reusing the native adaLN interface to transmit early-layer representations as modulation signals. This design naturally combines cross-layer history with the model's native condition, such as timestep and class embeddings. ResNorm allows the earlier structural representation to adaptively scale, shift, and gate deeper blocks. This not only produces more differentiated deep representations and reduces cross-layer similarity in later blocks, but also accelerates training convergence and improves performance. On ImageNet 256$\times$256, SiT-XL/2 + ResNorm achieves up to $\textbf{2}\times$ and $\textbf{36}\times$ faster training than SiT-XL/2 + REG and SiT-XL/2 + REPA. Moreover, SiT-L/2 + ResNorm trained for only 400K iterations already outperforms SiT-XL/2 + REG.

Current hierarchical attention methods, such as NSA and InfLLMv2, select the top-$k$ relevant key-value (KV) blocks based on coarse attention scores and subsequently apply fine-grained softmax attention on the selected tokens. However, the top-$k$ operation assumes the number of relevant tokens for any query is fixed and it precludes the gradient flow between the sparse and dense stages. In this work, we propose DashAttention (Differentiable and Adaptive Sparse Hierarchical Attention), which leverages the adaptively sparse $\alpha$-entmax transformation to select a variable number of blocks according to the current query in the first stage. This in turn provides a prior for the second-stage softmax attention, keeping the entire hierarchy fully differentiable. Contrary to other hierarchical attention methods, we show that DashAttention is non-dispersive, translating to better long-context modeling ability. Experiments with large language models (LLMs) show that DashAttention achieves comparable accuracy as full attention with 75% sparsity and a better Pareto frontier than NSA and InfLLMv2, especially in high-sparsity regimes. We also provide an efficient, GPU-aware implementation of DashAttention in Triton, which achieves a speedup of up to $3.3\times$ over FlashAttention-3 at inference time. Overall, DashAttention offers a cost-effective strategy to model long contexts.

Prefill-only KV compression freezes a token subset at the end of prefill and decodes from it without further eviction. The retention decision is therefore irreversible, yet existing methods estimate the corrective signals it relies on, per-head reliability and prompt-level compression sensitivity, online from a single noisy prompt. We argue this is the wrong statistical unit: these signals exhibit far higher cross-prompt regularity than within-prompt signal-to-noise. We introduce CompilerKV, a KV-retention policy whose corrective tables are compiled offline from a calibration corpus, reducing online correction after the standard observation-window scan to $O(1)$ lookups plus a budget clamp. We find that compiled retention tables behave as portable architectural priors: rankings transfer across disjoint corpora on four backbones, with mean Spearman $\bar{\rho}=0.90$, and direct model-to-model table transfer costs only $0.4$-$0.8$ LongBench points on average. At a 512-token budget, CompilerKV attains compressed-SOTA on all four backbones, improving over the strongest prefill-only baseline by $+1.67$ points on average, with task-bootstrap 95% CI $[+1.08,+2.37]$. Pressure regimes amplify the gap: under a fixed $512/32\mathrm{k}$ cache ratio, CompilerKV remains the strongest compressed method through 128k RULER, with approximately 73 versus FullKV 79 and SnapKV 38; on 32k NIAH it reaches $0.89$ versus SnapKV $0.42$; and at 32k input, retaining only $1.56\%$ of the prefill KV, batch-16 serving remains feasible where FullKV is OOM.

Selective state-space models (SSMs) have recently shown strong potential for efficient image restoration. A central difficulty, however, is that a 2D feature map typically needs to be linearized into a 1D token sequence for causal state-space updates, making restoration performance highly sensitive to the imposed scan order. Existing deterministic scans follow space-filling principles but remain fixed and content-agnostic, and can yield fragmented or semantically inconsistent contexts under degradations. We propose **CoScan**, a **co**ntent-adaptive **scan**-order design for causal vision SSMs. CoScan learns image-dependent space-filling traversals via differentiable proxy supervision driven by **scan-quality objectives** that encourage *short-range coherence and causal predictability, while preserving spatial adjacency*. To support hierarchical restoration backbones, we introduce an **efficient multi-scale scan construction** by *coarsening a fine-level minimum spanning tree*, producing coherent scans across feature resolutions. In addition, **a frequency-guided start selection strategy** is designed to *place informative structures early in the sequence prefix to strengthen causal context formation*. Plugged into existing Mamba-based backbones, CoScan yields consistent improvements across several restoration benchmarks with minimal overhead and shows promising results on high-level vision tasks. Code will be released upon acceptance.

Deep neural networks have achieved remarkable success across a variety of tasks, yet they often suffer from unreliable probability estimates, leading to overconfident predictions. Conformal Prediction (CP) offers a principled framework for uncertainty quantification, yielding prediction sets with rigorous coverage guarantees. Recent work has proposed integrating conformal objectives into training, optimizing for overall set size. Nevertheless, shaping the prediction sets in a class-conditional manner is not straightforward and typically requires prior knowledge of the data distribution. While class-conditional efficiency has been studied at the post-hoc calibration level (e.g., via class-conditional scoring or calibration partitioning), no prior conformal training method explicitly addresses it. To fill this gap, we introduce Class Adaptive Conformal Training (CaCT), which formulates conformal training as an augmented Lagrangian optimization problem that adaptively shapes prediction sets class-conditionally without relying on parametric distributional assumptions, beyond a user-specified global coverage constraint. Experiments on multiple benchmark datasets, including standard and long-tailed image recognition as well as text classification, demonstrate that CaCT consistently outperforms prior conformal training methods, reducing prediction set size while maintaining the desired coverage guarantees.

Adaptive learning requires distinguishing environmental volatility from observation stochasticity, two sources of uncertainty that demand opposite adjustments to the learning rate but inflate experienced variance similarly. Disentangling them is computationally difficult with no tractable closed-form solution. Particle-filter methods are the natural tool for this kind of joint inference, but their stochastic likelihoods and non-differentiable objectives force derivative-free fitting protocols and discourage the individual-difference analyses central to cognitive modeling, where small effect sizes leave little room for additional estimator noise. We introduce the Categorical Bayes Filter (CBF), a deterministic alternative that preserves the conditional structure of recent particle-filter accounts but replaces the stochastic outer layer with a categorical distribution on a quantile grid parameterized through differentiable Beta quantile functions. The procedure performs evidence maximization with an exact, deterministic marginal likelihood that is fully differentiable in the grid parameters. In a volatility-stochasticity task with N = 643 participants, fitted CBF dispersion parameters reveal a cross-over phenotyping pattern between volatility-blind and stochasticity-blind subjects that is not recoverable from particle-filter parameters fit to the same data under a state-of-the-art protocol. The deterministic structure also yields a trial-by-trial ambiguity signal that predicts response times not used in fitting. More broadly, the approach opens individual-level analyses in cognitive modeling and computational psychiatry that stochastic methods have effectively foreclosed.