Post-training pretrained autoregressive models (ARMs) into masked diffusion models (MDMs) provides an efficient route to diffusion language modeling, but it remains unclear whether the resulting models reuse inherited autoregressive computation or reorganize it for non-autoregressive generation. We compare two 7B ARM-MDM families across four controlled diagnostic tasks and find a task-dependent mechanism shift. On prefix-dominant tasks, MDMs largely preserve inherited high-attribution pathways or exhibit only modest changes in where computation occurs. On globally constrained tasks, the reorganization is substantially stronger, with task-relevant computation shifting toward earlier layers. This depth-wise pattern persists across prompt resampling, circuit budgets, and tested inference budgets, while targeted ablations support the functional importance of the identified structures under the tested intervention protocols. At the component level, diagnostic probes suggest that ARMs rely more strongly on sharply specialized components, whereas MDMs exhibit weaker single-component specialization and more diffuse output-space alignment. Together, these results suggest that diffusion post-training selectively preserves or reorganizes inherited computation according to task structure, rather than uniformly replacing autoregressive mechanisms.
Personalized Federated Learning (PFL) excels at tailoring client-specific models, which is particularly critical for decentralized and heterogeneous data environments, yet existing methods produce static models with a fixed tradeoff between accuracy and efficiency. This inherent static nature limits their ability to adapt to inference demands that vary with context and resource availability, posing a challenge for real-world deployment. Early-exit networks (EENs), which enable adaptive inference via intermediate classifiers, offer a promising solution. However, integrating EENs into PFL introduces two intertwined conflicts: client-wise heterogeneity across clients and depth-wise interference arising from conflicting exit objectives. Prior studies fail to resolve both conflicts simultaneously, leading to suboptimal performance. In this paper, we propose X-FED, a novel Conflict-Aware Cross-Client Federated Exit Distillation framework that jointly addresses both client- and depth-wise conflicts while extending PFL to early-exit networks. At its core, X-FED employs a progressive, depth-prioritized student coordination mechanism that mitigates interference among shallow and deep exits while enabling effective personalized knowledge transfer across clients. Furthermore, we introduce a client-decoupled formulation that reduces communication overhead with theoretical soundness. Extensive evaluations on various datasets show that compared to state-of-the-art PFL and PFL-EE methods, X-FED achieves higher accuracy while reducing inference costs by 30.79%-46.86%.
Boyi Liu, Zimu Zhou, Cheng Fang et al.· 0 citations
Large Language Models (LLMs) exhibit systematic biases across demographic groups. Auditing is proposed as an accountability tool for black-box LLM applications, but suffers from resource-intensive query access. We conceptualise auditing as uncertainty estimation over a target fairness metric and introduce BAFA, the Bounded Active Fairness Auditor for query-efficient auditing of black-box LLMs. BAFA maintains a version space of surrogate models consistent with queried scores and computes uncertainty intervals for fairness metrics (e.g., $\Delta$ AUC) via constrained empirical risk minimisation. Active query selection narrows these intervals to reduce estimation error. We evaluate BAFA on two standard fairness dataset case studies: \textsc{CivilComments} and \textsc{Bias-in-Bios}, comparing against stratified sampling, power sampling, and ablations. BAFA achieves target error thresholds with up to 40$\times$ fewer queries than stratified sampling (e.g., 144 vs 5,956 queries at $\varepsilon=0.02$ for \textsc{CivilComments}) for tight thresholds, demonstrates substantially better performance over time, and shows lower variance across runs. These results suggest that active sampling can reduce resources needed for independent fairness auditing with LLMs, supporting continuous model evaluations.
David Hartmann, Lena Pohlmann, Lelia Hanslik et al.· 0 citations
Long-context LLM inference is bottlenecked by the quadratic attention complexity and linear Key-Value (KV) cache growth. Prior approaches mitigate this via post-hoc selection or eviction but overlook the root inefficiency: indiscriminate token admission. In this paper, we formalize KV management as a causal system of three primitives: KV Admission, Selection, and Eviction. We instantiate KV Admission via Write-Gated KV (WG-KV), a lightweight mechanism that learns to predict token utility before cache entry. By filtering out redundant states early to maintain a compact global cache alongside a sliding local cache, WG-KV significantly reduces memory usage and accelerates both prefill and decode phases. Our results demonstrate that learning what to write is a principled and practical recipe for efficient long-context inference. Code is available at https://github.com/EMCLab-Sinica/WG-KV.
Yen-Chieh Huang, Pi-Cheng Hsiu, Rui Fang et al.· 0 citations
Reach audiences
Advertise in front of researchers, engineers, and readers.
Attention is the dominant source of latency during long-context LLM inference, an increasingly popular workload with reasoning models and RAG. We propose Kascade, a training-free sparse attention method that leverages known observations such as 1) post-softmax attention is intrinsically sparse, and 2) the identity of high-weight keys is stable across nearby layers. Kascade computes exact Top-k indices in a small set of anchor layers, then reuses those indices in intermediate reuse layers. The anchor layers are selected algorithmically, via a dynamic-programming objective that maximizes cross-layer similarity over a development set, allowing easy deployment across models. The method incorporates efficient implementation constraints (e.g. tile-level operations), across both prefill and decode attention. The Top-k selection and reuse in Kascade is head-aware and we show in our experiments that this is critical for high accuracy. Kascade achieves up to 4.1x speedup in decode attention and 2.2x speedup in prefill attention over FlashAttention-3 baseline on H100 GPUs while closely matching dense attention accuracy on long-context benchmarks such as LongBench and AIME-24.
Dhruv Deshmukh, Saurabh Goyal, Nipun Kwatra et al.· 0 citations
We study how explicit world-modeling objectives affect the internal representations and downstream capability of Transformers, using Rubik's Cubes as our training domain. We ask: (1) how does explicitly pretraining a world model affect a model's latent representations, (2) how does world-model quality affect post-training performance, and (3) how should a finite data budget be split between pretraining and task-specific fine-tuning? We compare standard action-prediction fine-tuning with two strategies that add explicit state-prediction supervision: state pretraining followed by fine-tuning, and joint action and state training. We measure task accuracy after Group Relative Policy Optimization (GRPO). We further find that explicit world-modeling yields better representations in terms of higher probing accuracy and steerability of the model, and that better representations yield larger gains from GRPO, especially on harder cube states. Finally, when the fine-tuning budget is held fixed, probe accuracy strongly predicts the GRPO improvement. Under a fixed total data budget, accuracy is maximized by allocating only a small fraction to pretraining.
Prakhar Gupta, Henry Conklin, Sarah-Jane Leslie et al.· 0 citations
Training process reward models (PRMs) requires step-level correctness labels, obtained either through expensive human annotation or by relying on ground-truth answers, limiting the ability to scale process-level supervision. We propose ScalePRM, which scales verification compute as an alternative: given a problem and a candidate solution, we generate multiple independent verifications of each reasoning step and aggregate their judgments to produce synthetic step-level labels without ground truth. We explore two representative inference-time scaling strategies, parallel scaling through self-consistency and sequential scaling through meta-critique, and train generative PRMs on the resulting synthetic data. On ProcessBench, a benchmark for identifying erroneous steps in mathematical reasoning, PRMs trained on step-level self-consistency data achieve 67.5 F1, surpassing reference-guided training with ground-truth access (66.4 F1) and GPT-4o as a critic (61.9 F1). When deployed as reward signals in RL training with Qwen2.5-Math-7B, our best PRM achieves 47.4% average accuracy across six mathematical reasoning benchmarks, outperforming ground-truth-based RLVR (43.9%). We also identify and address reward exploitation patterns unique to generative PRM-based RL. Our results demonstrate that scaling verification compute is a viable alternative to ground-truth supervision for training process reward models.
Salman Rahman, Sruthi Gorantla, Arpit Gupta et al.· 0 citations
Growing demands from tasks like code generation, deep reasoning, and long-document understanding have made long-context generation a crucial capability for large language models (LLMs). Speculative decoding is one of the most direct and effective approaches for accelerating generation. It follows a draft-verify paradigm, where a lightweight draft model proposes several candidate tokens and the target model verifies them. However, we find that as the context length grows, verification becomes the dominant bottleneck. To further accelerate speculative decoding in long-context generation, we introduce SpecPV, a self-speculative decoding approach that performs fast verification using partial key-value states (KV) and periodically applies full verification to eliminate accumulated errors. We validate SpecPV across multiple long-context benchmarks and models, including LLaMA-3.1-8B-Instruct and Qwen3-series. Experimental results show that SpecPV achieves up to 6x decoding speedup over standard autoregressive decoding with minor degradation.
Zhendong Tan, Xingjun Zhang, Chaoyi Hu et al.· 0 citations
Real-world machine learning (ML) pipelines rarely produce a single model; instead, they produce a Rashomon set of many near-optimal ones. We show that this multiplicity reshapes key aspects of trustworthiness. At the individual-model level, sparse interpretable models tend to preserve privacy but are fragile to adversarial attacks. In contrast, the diversity within a large Rashomon set enables reactive robustness: even when an attack compromises one model, a practitioner can switch to a different near-optimal model that remains accurate, without retraining. However, the same diversity increases information leakage, as disclosing more near-optimal models provides an attacker with progressively richer views of the training data. This produces a robustness-privacy trade-off governed by diversity, which we analyze theoretically and empirically. Beyond this trade-off, Rashomon sets are stable under small distribution shifts, so a set computed once remains valid under such shifts without re-computation. Our results highlight the dual role of Rashomon sets as both a resource and a risk for trustworthy ML.
Ethan Hsu, Harry Chen, Chudi Zhong et al.· 0 citations
Deploying deep neural networks on resource-constrained devices faces two critical challenges: maintaining accuracy under aggressive quantization while ensuring predictable inference latency. We present a curiosity-driven quantized Mixture-of-Experts framework that addresses both through Bayesian epistemic uncertainty-based routing across heterogeneous experts (BitNet ternary, 1-16 bit BitLinear, post-training quantization). Evaluated on audio classification benchmarks (ESC-50, Quinn, UrbanSound8K), our 4-bit quantization maintains 99.9 percent of full-precision F1 (0.858 vs 0.859) with 4x compression and 31 percent energy savings versus 8-bit, while both achieve statistical parity with full precision (p > 0.05).
Crucially, curiosity-driven routing simultaneously improves accuracy and stability: on Quinn, F1 increases from 0.802 to 0.809 while cross-fold variance drops by 85 percent (p < 0.001, Levene's test), with reductions of 50 to 94 percent across datasets. The routing is self-organizing, with the high-precision 8-bit expert automatically receiving the most uncertain samples (20 percent lower confidence, p < 0.001), while lightweight experts handle easier inputs. Datasets with already low baseline variance show no artificial stability gain, confirming the mechanism targets genuine epistemic uncertainty rather than overfitting routing decisions.
At 1.2M parameters, the framework provides interpretable, precision-aware routing suitable for safety-sensitive edge deployments where both accuracy and predictability are critical.
Sebasti\'an Andr\'es Cajas Ord\'o\~nez, Luis Fernando Torres Torres, Mackenzie J. Meni et al.· 0 citations
This paper focuses on dynamic origin-destination matrix estimation (DODE), a crucial calibration process necessary for the effective application of microscopic traffic simulations. The fundamental challenge of the DODE problem in microscopic simulations stems from the complex temporal dynamics and inherent uncertainty of individual vehicle dynamics. This makes it highly challenging to precisely determine which vehicle traverses which link at any given moment, resulting in intricate and often ambiguous relationships between origin-destination (OD) matrices and their contributions to resultant link flows. This phenomenon constitutes the credit assignment problem, a central challenge addressed in this study. We formulate the DODE problem as a Markov Decision Process (MDP) and propose a novel framework that applies model-free deep reinforcement learning (DRL). Within our proposed framework, the agent learns an optimal policy to sequentially generate OD matrices, refining its strategy through direct interaction with the simulation environment. This approach was evaluated through a toy experiment on the Nguyen-Dupuis network and a case study utilizing an actual highway subnetwork spanning Santa Clara and San Jose. Experimental results show that the proposed method consistently improves calibration performance relative to the strongest conventional baseline, reducing link-flow MSE by 23.7% in the toy experiment and by 59.2-88.3% in the real-world case study. By reframing DODE as a sequential decision-making problem, our approach addresses the credit assignment challenge through a learned policy and provides a novel framework for calibration of microscopic traffic simulations.
Physics-Informed Neural Networks (PINNs) offer a flexible framework for solving nonlinear partial differential equations (PDEs), yet conventional implementations often fail to preserve key physical invariants during long-term integration. This paper introduces a \emph{structure-preserving PINN} framework for the nonlinear Korteweg--de Vries (KdV) equation, a prototypical model for nonlinear and dispersive wave propagation. The proposed method embeds the conservation of mass and Hamiltonian energy directly into the loss function, ensuring physically consistent and energy-stable evolution throughout training and prediction. Unlike standard \texttt{tanh}-based PINNs~\cite{raissi2019pinn,wang2022modifiedpinn}, our approach employs sinusoidal activation functions that enhance spectral expressiveness and accurately capture the oscillatory and dispersive nature of KdV solitons. Through representative case studies -- including single-soliton propagation (shape-preserving translation), two-soliton interaction (elastic collision with phase shift), and cosine-pulse initialization (nonlinear dispersive breakup) -- the model successfully reproduces hallmark behaviors of KdV dynamics while maintaining conserved invariants. Ablation studies demonstrate that combining invariant-constrained optimization with sinusoidal feature mappings accelerates convergence, improves long-term stability, and mitigates drift without multi-stage pretraining. These results highlight that computationally efficient, invariant-aware regularization coupled with sinusoidal representations yields robust, energy-consistent PINNs for Hamiltonian partial differential equations such as the KdV equation.
A new machine-learning framework aims to improve the success rate of computational protein design while moving away from results that reproduce sequences found in nature.
MIT News · Artificial Intelligence· news.mit.eduAug 24, 2026
A new method for surgically removing training examples from a model reveals that as datasets grow, the link between what a model learns and what it produces dissolves.
