Probing Routing-Conditional Calibration in Attention-Residual Transformers

2026-05-10

AI Summary

This study investigates whether routing pathways in attention residual Transformers can provide stable calibration signals beyond model confidence. To this end, we develop a diagnostic framework that jointly controls for matched confidence, bandwidth, and capacity, and introduce— for the first time—group permutation tests and multi-class control probes. We systematically evaluate the authenticity of routing-conditioned calibration using Nadaraya-Watson kernel regression, full-vector MLPs, and multiple bandwidth selection strategies, including Scott’s rule-of-thumb multiplier, cross-validated negative log-likelihood (CV-NLL), and global expected calibration error (ECE). Our experiments reveal that routing-derived summary features do not yield consistent calibration gains, and the apparent advantages of sophisticated routing-aware probes vanish under appropriate controls, indicating that routing information in current architectures does not reliably enhance calibration performance.

Abstract

Post-hoc calibration is usually evaluated as a function of logits or softmax confidence alone, even as routing-augmented architectures increasingly accompany predictions with sample-specific internal routing traces and pair them with claims of calibration-relevant uncertainty. We ask a basic question: do these traces provide stable routing-specific evidence for post-hoc calibration beyond confidence? We study this in Attention-Residual transformers (Kimi Team, 2026) through a matched-confidence diagnostic suite that stratifies examples by routing-derived state, compares subgroup gaps against within-bin routing-permutation nulls, and evaluates matched post-hoc probes differing only in their auxiliary feature. Across our completed AR runs, scalar routing summaries do not provide stable evidence of routing-conditional miscalibration: weighted gaps remain small or seed-sensitive, and only $1$ of $30$ within-bin permutation tests rejects the conditional-null at $α=0.05$ (only on one seed; not stable across seeds in that cell). AR-CondCal, a minimal $2$-D Nadaraya--Watson probe on confidence and routing-depth variance, lies within the seed-variance band of matched confidence-only and predictive-entropy controls and does not reliably improve worst-routing-tertile ECE; bandwidth-sensitivity checks (Scott multiples, CV-NLL, global-ECE oracle) do not change this. A full-vector MLP over $(c, H_1, \ldots, H_L)$ can appear to improve over a linear confidence baseline, but the apparent gain disappears once a capacity-matched confidence-only MLP is included as a control, and shuffled routing profiles achieve comparable performance. Apparent routing-aware calibration gains in this AR setting should not be read as internal-state calibration until matched-confidence, bandwidth, capacity, and permutation controls rule out common confounds.

Problem

• post-hoc calibration
• routing traces
• conditional calibration
• Attention-Residual Transformers
• calibration evidence

Innovation

• routing-conditional calibration
• post-hoc calibration
• Attention-Residual Transformers
• permutation null tests
• matched-confidence diagnostics