CAM: Uncertainty-Aware Transformer Attention

• Credal Attention Mechanism (CAM) is an uncertainty-aware reformulation of Transformer attention that uses a Dirichlet distribution to express epistemic uncertainty.
• CAM computes attention by converting raw scores into non-negative evidence, producing Dirichlet parameters that form a credal set of plausible attention distributions.
• The mechanism offers practical benefits like improved out-of-distribution detection, reduced hallucinations, and a drop-in replacement for standard softmax attention.

Searching arXiv for the specified paper to ground the article in the primary source. arxiv_search(query="(Ji et al., 14 Oct 2025)", max_results=5, sort_by="submittedDate") The Credal Attention Mechanism (CAM) is an evidential reformulation of Transformer attention introduced in the Credal Transformer, where standard softmax-based attention is replaced by a credal representation that models uncertainty in the attention distribution itself. Rather than collapsing attention scores into a single categorical distribution, CAM constructs a credal set—represented implicitly by a Dirichlet distribution over attention vectors—so that ambiguity in token-to-token relevance remains explicit throughout the attention computation. In this framework, the diffuseness of the Dirichlet, summarized by its total concentration and derived vacuity, serves as a direct measure of epistemic uncertainty. The mechanism is proposed as a response to Artificial Certainty, the claim that softmax forces a sharp decision even when the underlying evidence is weak or conflicting, thereby contributing to hallucinations and overconfident errors in LLMs (Ji et al., 14 Oct 2025).

1. Conceptual basis and motivation

CAM is motivated by the observation that standard Transformer attention converts raw compatibility scores into a single normalized distribution via softmax. For a query vector $\mathbf{Q}_i$ and keys $\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$, the standard formulation computes attention scores $s_{ij}$ and then applies

$$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$

This operation yields one sharp probability distribution even when the score pattern is ambiguous. The stated consequence is that the model is forced to “pick a winner” at every layer, propagating overconfident intermediate states through the network (Ji et al., 14 Oct 2025).

The central claim is that this induces Artificial Certainty: ambiguous intermediate representations are collapsed into overly confident attention distributions, and the repeated composition of these distributions can contribute to hallucinations, particularly under out-of-distribution conditions or for unanswerable questions. CAM addresses this by replacing the single attention vector with a family of plausible attention distributions. A small/sharp credal set corresponds to high certainty, whereas a large/diffuse credal set corresponds to epistemic uncertainty. This suggests that CAM is not merely a calibration device at the output layer, but an attempt to internalize uncertainty within the core inference mechanism of the Transformer (Ji et al., 14 Oct 2025).

2. Evidential formulation and credal semantics

CAM draws on the evidence-based uncertainty perspective associated with Evidential Deep Learning. In that setting, a model produces non-negative evidence values $e_j \ge 0$ that are converted into Dirichlet concentration parameters according to

$\alpha_j = e_j + 1.$

The resulting Dirichlet distribution over categorical probability vectors $\mathbf{p} = (p_1,\dots,p_L)$ is

$$\text{Dir}(\mathbf{p} \mid \boldsymbol{\alpha}) = \frac{1}{B(\boldsymbol{\alpha})} \prod_{j=1}^{L} p_j^{\alpha_j - 1},$$

with total concentration

$\alpha_0 = \sum_{j=1}^{L} \alpha_j.$

Large concentration values yield a peaked Dirichlet and therefore high certainty, while small concentrations yield a flat or diffuse Dirichlet and therefore high uncertainty (Ji et al., 14 Oct 2025).

Within CAM, the Dirichlet is interpreted as inducing a credal set, that is, a family of plausible categorical attention distributions. The paper does not provide a formal set-theoretic derivation of belief and plausibility functions, but it adopts the core evidential idea that uncertainty should be represented as a set of distributions rather than a single point estimate. For each query position $i$, CAM defines a Dirichlet $\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$0, and the distributions plausible under this Dirichlet constitute the position-wise credal set. A plausible implication is that CAM operationalizes epistemic uncertainty in attention without requiring explicit combinatorial manipulation of sets of probabilities, since the Dirichlet provides a tractable parametric surrogate (Ji et al., 14 Oct 2025).

3. Mathematical definition of the mechanism

CAM preserves the standard attention score computation but changes its interpretation. Instead of treating attention scores as logits for a categorical distribution, it treats them as sources of evidence. For each query position $\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$1 and key position $\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$2,

$\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$3

$\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$4

$\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$5

The exponential ensures non-negative evidence, and the additive constant ensures $\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$6, so that even zero evidence yields a valid concentration parameter corresponding to a uniform prior (Ji et al., 14 Oct 2025).

Let

$\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$7

CAM does not sample an attention vector from the Dirichlet. Instead, for computational tractability it uses the Dirichlet mean as the effective attention distribution: $\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$8 Substituting the definition of $\{\mathbf{K}_1,\dots,\mathbf{K}_L\}$9 yields

$s_{ij}$0

This is not identical to softmax: it incorporates both the evidence magnitude and the prior offset induced by the $s_{ij}$1 term. The paper notes that when evidence is high and differentiated, standard attention is effectively recovered because

$s_{ij}$2

which is close to the softmax over $s_{ij}$3 (Ji et al., 14 Oct 2025).

The uncertainty statistic used by CAM is the vacuity

$s_{ij}$4

where $s_{ij}$5 is the number of attention targets. High total evidence drives $s_{ij}$6 toward $s_{ij}$7, while low evidence gives $s_{ij}$8 and thus $s_{ij}$9, representing maximal vacuity. The paper emphasizes that this quantity is directly differentiable with respect to the attention scores and functions as a local, per-head per-position measure of epistemic uncertainty (Ji et al., 14 Oct 2025).

4. Integration into Transformer attention

In single-head form, CAM substitutes a Dirichlet-based normalization pipeline for softmax normalization. Given

$$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$0

the mechanism computes

$$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$1

$$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$2

$$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$3

and then forms the attention output as

$$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$4

Algorithmically, this preserves the standard structure of attention: the score matrix is still built from queries and keys, the resulting normalized weights are still row-wise, and the value aggregation step is unchanged. The difference lies entirely in how the normalized weights are produced and in the additional uncertainty information retained through $$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$5 and $$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$6 (Ji et al., 14 Oct 2025).

The multi-head architecture is not otherwise altered. CAM is applied head-wise in place of softmax attention; each head computes its own Dirichlet parameters and vacuity values, the head outputs are concatenated, and the usual linear transformation follows. The paper states that uncertainty signals can be propagated or aggregated, but the reported usage is primarily at the final layer, particularly for abstention decisions. This suggests that CAM is intended as a drop-in replacement at the attention sub-block level rather than a redesign of the broader Transformer stack (Ji et al., 14 Oct 2025).

The mechanism is described as supporting a continuous transition between confident and diffuse attention. When scores are large and sharply differentiated, the exponentiated evidence terms dominate the additive prior, the total concentration becomes large, vacuity becomes small, and the resulting attention vector is peaked. When scores are small or similar, the concentration parameters remain near $$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$7, total concentration stays close to $$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$8, vacuity becomes high, and the effective attention vector becomes diffuse. Crucially, the model retains an explicit signal indicating that this diffuse pattern reflects uncertainty rather than mere averaging (Ji et al., 14 Oct 2025).

5. Uncertainty quantification, abstention, and reliability claims

The principal uncertainty measure associated with CAM is the attention-level vacuity $$\mathbf{a}_i = \text{Softmax}(\mathbf{s}_i), \quad a_{ij} = \frac{\exp(s_{ij})}{\sum_{k=1}^{L} \exp(s_{ik})}.$$9, computed for each query position and each head. The paper treats this as a scalar measure of epistemic uncertainty over where to attend. In practical use, this uncertainty is linked to changes in model behavior rather than being only diagnostic. Specifically, high uncertainty can be used as a cue to abstain, to flag ambiguity, or to identify inputs that are out of distribution (Ji et al., 14 Oct 2025).

For out-of-distribution detection, a classifier built from a Credal Transformer encoder was trained only on in-distribution data and then evaluated on in-distribution, out-of-distribution, and pure noise inputs. The reported average uncertainty at the final layer was 0.0415 for in-distribution data, 0.1378 for out-of-distribution data, and 0.1953 for nonsense inputs. The reported pattern is therefore monotonic: uncertainty increases as inputs move farther from the training distribution. The intended contrast is with standard Transformers, which are described as often producing high-confidence predictions on out-of-distribution and nonsense inputs (Ji et al., 14 Oct 2025).

On ambiguous inputs, CAM is reported to produce larger credal sets and higher $e_j \ge 0$0, indicating that multiple attention patterns remain tenable. Standard attention, by contrast, is described as being forced to commit to a single normalized distribution. A plausible implication is that CAM can distinguish between uncertainty caused by lack of evidence and confidence produced by a normalization artifact, although the paper frames this primarily qualitatively rather than through a formal decomposition of uncertainty types (Ji et al., 14 Oct 2025).

For question answering with unanswerable questions, the Credal Transformer uses high uncertainty—either through $e_j \ge 0$1 or an analogous output-layer uncertainty measure—as a cue to abstain rather than fabricate an answer. The paper reports a significant reduction in confident errors relative to standard LLMs, while not specifying exact numerical gains in the provided excerpt. The abstention rule is described conceptually as threshold-based: if uncertainty exceeds a threshold, the system abstains or flags the answer as unknown; otherwise it proceeds with normal prediction (Ji et al., 14 Oct 2025).

6. Optimization, complexity, and implementation properties

The Credal Transformer paper presents CAM as fully differentiable and trainable with standard optimization techniques. The model adopts the evidential principle of outputting evidence, but the paper does not provide an explicit CAM-specific loss or the full Evidential Deep Learning loss machinery such as KL regularization to a prior. Instead, standard supervised objectives, such as cross-entropy on final outputs, can be used while gradients propagate through

$e_j \ge 0$2

The uncertainty $e_j \ge 0$3 is described as being used primarily at evaluation or decision time rather than being directly optimized (Ji et al., 14 Oct 2025).

The architectural integration cost is presented as small. CAM is described as a drop-in replacement for the attention module: softmax is replaced with the mapping $e_j \ge 0$4, while the rest of the Transformer layer—multi-head structure, feed-forward blocks, residual connections, and layer normalization—remains unchanged. The reported empirical complexity values are summarized below.

Quantity	Standard attention	CAM
GFLOPs	25.77 G	25.77 G (+0%)
Inference time	—	+4.4%
Training step time	—	+11.6%

These figures support the claim that CAM achieves uncertainty quantification with negligible computational cost in the reported setting. This suggests that the mechanism alters the semantics of normalization more than the asymptotic structure of attention, although the paper does not provide a separate big-$e_j \ge 0$5 derivation (Ji et al., 14 Oct 2025).

7. Interpretation, limitations, and prospective extensions

CAM is presented as a principled architectural attempt to encode epistemic uncertainty directly within the attention mechanism. In that interpretation, the contribution is not solely a method for producing abstention signals, but a shift from viewing attention as a deterministic probability assignment to viewing it as an uncertainty-aware evidential object. The associated broader claim is that hallucination mitigation can be approached by redesigning internal inference so that the model can “know when it doesn’t know,” rather than depending only on post-hoc confidence heuristics (Ji et al., 14 Oct 2025).

The paper also identifies several limitations. First, the reported empirical coverage is concentrated on relatively controlled discriminative settings, including out-of-distribution detection and question answering with abstention, while effectiveness in open-ended long-form generation remains insufficiently characterized. Second, the uncertainty signal $e_j \ge 0$6 is mostly used at the output or decision level; richer internal uses such as dynamically reweighting attention heads or modulating information flow based on certainty are proposed but not implemented. Third, although the reported overhead is small on the evaluated benchmarks, scaling CAM to 100B+ parameter LLMs under distributed training is described as an open issue, with numerical stability and memory overhead for storing uncertainty noted as practical concerns (Ji et al., 14 Oct 2025).

Several future directions are explicitly suggested: integrating CAM uncertainty into decoding strategies such as refusal behavior or dynamic top-$e_j \ge 0$7/temperature control; extending credal mechanisms beyond attention to other components such as feed-forward layers; exploring head-level and layer-level uncertainty aggregation for interpretability and control; and evaluating larger-scale systems on more diverse and safety-critical benchmarks. These proposals indicate that CAM is framed not only as a specific attention replacement, but as the beginning of a broader uncertainty-aware design paradigm for Transformer architectures (Ji et al., 14 Oct 2025).