Lévy Attention: Single-Pass Predictive Uncertainty for Continuous-Time Attention

Lévy Attention: Single-Pass Predictive Uncertainty for Continuous-Time Attention§

Core Idea§

Lévy Attention replaces the standard softmax attention with a stochastic operator whose output is a random integral against an inhomogeneous Poisson random measure. Query-key compatibilities define an intensity function over a continuous (time × channel) space, and the attention output is the average of a value field at random atoms sampled from this measure. In expectation, this reduces to a deterministic mollified cosine-kernel attention, so the stochastic operator can replace a softmax layer while still training with exact gradients. Crucially, the Poisson formulation preserves quantities that softmax discards: the total compatibility mass (evidence $\Lambda_q$) and the value spread ($\mathrm{tr}\,\Sigma_V(q)$). An exact variance identity combines these into a closed-form root-mean-square deviation (RMSD) $\hat{\sigma}(q)$, giving per-query predictive uncertainty in a single deterministic forward pass.

Mathematical Formulation§

For a query $q$ at continuous timestamp $t$, keys $k_i$ and values $v_i$ at irregular timestamps, the attention output is:

$$ \mathrm{L\acute{e}vyAtt}(q) = \frac{1}{\Lambda_q} \int_{\mathcal{X}} v(x) \, \Pi(dx) $$

where $\Pi$ is a Poisson random measure on index space $\mathcal{X}$ with intensity $\lambda_q(x) = \exp(\mathrm{sim}(q, x))$ (cosine similarity scaled), $\Lambda_q = \int \lambda_q(x) dx$ is the total evidence, and $v(x)$ interpolates the value field. The conditional expectation and variance are:

$$ \mathbb{E}[\mathrm{L\acute{e}vyAtt}(q) \mid \Lambda_q] = \frac{1}{\Lambda_q} \int \lambda_q(x) v(x) dx $$

$$ \mathrm{Var}(\mathrm{L\acute{e}vyAtt}(q) \mid \Lambda_q) = \frac{1}{\Lambda_q} \int \lambda_q(x) \|v(x) - \bar{v}\|^2 dx $$

Define $\Sigma_V(q)$ as the value covariance matrix over the interpolated value field, and $\varphi(\Lambda_q)$ as a known function (e.g., $1/\Lambda_q$ for unit-rate Poisson). The closed-form uncertainty is:

$$ \hat{\sigma}(q) = \sqrt{\mathrm{tr}\,\Sigma_V(q) \, \varphi(\Lambda_q)} $$

This is the RMSD of the sampled operator, computed deterministically from the same quantities used to produce the mean output — no extra network head or sampling required.

Implementation Details§

• Intensity construction: Queries and keys are embedded via a trainable function into a continuous index space. The compatibility is a scaled cosine similarity, exponentiated to form a non-negative intensity.
• Value interpolation: Values are represented as a continuous field (e.g., via linear or Gaussian interpolation) over the index space, evaluated at Poisson atoms during sampling.
• Training: The mean operation is differentiable; gradients are computed exactly using the closed-form expectation, avoiding REINFORCE-type estimators. The variance identity is used only for uncertainty estimation at inference, but can also be used as a regularizer.
• Complexity: The deterministic pass costs the same as standard attention; the Poisson sampling for Monte Carlo estimation is only needed if one wants empirical samples (e.g., for comparison), but the closed-form $\hat{\sigma}$ is emitted at zero extra cost.

Code Sketch (PyTorch-style)§

import torch
import torch.nn as nn

class LevyAttention(nn.Module):
    def __init__(self, dim, num_heads):
        super().__init__()
        self.scale = dim ** -0.5
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)

    def forward(self, q, k, v, timestamps_q, timestamps_k):
        Q = self.q_proj(q) * self.scale
        K = self.k_proj(k)
        V = self.v_proj(v)

        # Intensity: exp(cosine(q,k)) over continuous time (approximated on grid)
        sim = torch.matmul(Q, K.transpose(-2, -1))  # [B, H, Q_len, K_len]
        Lambda = sim.exp().sum(dim=-1, keepdim=True)  # evidence mass

        # Expected output (mollified attention)
        attn = sim.softmax(dim=-1)
        mean_out = torch.matmul(attn, V)

        # Disagreement / value spread
        V_bar = mean_out.unsqueeze(-1)
        diff = V.unsqueeze(-2) - V_bar  # [B, H, Q_len, K_len, dim]
        Sigma_tr = (attn.unsqueeze(-1) * diff.pow(2).sum(-1)).sum(-2, keepdim=True)

        # Closed-form uncertainty: tr(Sigma) * phi(Lambda), with phi(Lambda)=1/Lambda
        sigma_hat = torch.sqrt(Sigma_tr / Lambda)

        return mean_out, sigma_hat  # sigma_hat is per-query RMSD

Empirical Findings§

• Uncertainty quality: The disagreement signal ($\mathrm{tr}\,\Sigma_V$) carries most predictive power, while the evidence factor ($\Lambda_q$) becomes informative only on sparse data (where softmax's normalization discards crucial magnitude information).
• On t-PatchGNN benchmarks, swapping the attention operator costs at most 5.6% accuracy on dense data and nothing on the sparsest dataset, while providing free uncertainty.
• The free single-pass $\hat{\sigma}$ outperforms 20-pass MC Dropout on matched five-seed suites in calibration and sharpness.
• A scaled Gaussian predictive with $\hat{\sigma}$ achieves lower CRPS (continuous ranked probability score) than a 50-draw sampler.
• A split-conformal wrapper around $\hat{\sigma}$ achieves nominal coverage at every confidence level, and ranking 3,383 unseen patients by trust takes only 1.4 seconds in one pass.

Significance§

This work offers a principled way to endow any cross-attention layer with built-in predictive uncertainty without architectural changes beyond the attention formulation. It leverages stochastic process theory (Poisson random measures) to recover information discarded by softmax normalization, enabling reliable uncertainty estimates for continuous-time forecasting and survival analysis in medical and sensor applications.