A lightweight controller let diffusion language agents decide how long to think, cutting inference FLOPs 20.6 percent on average and up to 77 percent against one baseline
AspireAI community · arXiv
*Diffusion language models can generate tokens in parallel, but an agent talking to an environment still needs to decide where to spend computation and how much text to produce before it is worth stopping.* Metacognitive Decoding adds a lightweight plug-in controller to the diffusion backbone, trained on agent trajectories, that learns when to finalise tokens and how far generation should extend as a function of what the interaction actually needs. Across five benchmarks it reduces inference FLOPs by 20.6 percent on average against the most efficient baseline while keeping competitive accuracy, and against Fast-dLLM the reduction runs from 56.6 to 77.2 percent, with accuracy gains of 28.20 and 20.96 points on the BFCL Non-live and Live tool-call benchmarks over LLaDA-8B-Tools. The ablations separate the two halves of the controller - disabling the flux decoding policy adds 24.5 percent FLOPs and disabling the confidence alignment policy adds 16.3 percent - and the multi-hop QA results are the strongest case, with gains of 9.50, 17.50 and 2.50 points on HotpotQA, 2Wiki and MuSiQue alongside an 88.3 to 89.7 percent FLOPs cut. The limitation is that the accuracy movement is not uniform, with MuSiQue actually giving up 0.40 points, so the efficiency win is much better established than the quality one.
Read the story and explore more from the AspireAI community.
Two sparse-autoencoder analyses have resisted scaling: matching semantically similar features across layers, and compressing large feature circuits into interpretable supernodes. This paper argues both are the same underlying problem, estimating semantic distances between SAE features that sit on different activation manifolds. Its representation treats each feature not as a single decoder vector but as an activation-weighted distribution over the hidden states expressing it, projects those distributions into a shared reference space, and compares them with Wasserstein distance. The representation is proved invariant to activation rescaling and stable under perturbation, and recovers true matches under finite-sample margin conditions. Empirically it outperformed decoder-vector and LLM-based baselines and captured subtle functional distinctions between related features, and it compressed large feature circuits into interpretable supernodes automatically. The limitation is that this is interpretability tooling: the paper demonstrates better feature matching and circuit compression, not that downstream interpretability claims become more accurate.