Example: Transformer-Style Attention Block

Attention stacks the operations that stress an abstract domain: two products of abstract operands (q @ k.T, attn @ v) with a softmax in between. The same script runs under the zonotope and the sparse polynomial domains — only the input class and the interpret object change.

import math

import torch
from torch import nn

from boundlab import Error
from boundlab.utils import ShapeDtype
from boundlab.zono import Zono
from boundlab.zono import interpret as zono_interpret
from boundlab.polysp import PolySp
from boundlab.polysp import interpret as polysp_interpret


class TinyAttention(nn.Module):
    def __init__(self, d_model=3):
        super().__init__()
        self.W_Q = nn.Linear(d_model, d_model)
        self.W_K = nn.Linear(d_model, d_model)
        self.W_V = nn.Linear(d_model, d_model)
        self.scale = math.sqrt(d_model)

    def forward(self, x):
        # x: [seq_len, d_model]
        q, k, v = self.W_Q(x), self.W_K(x), self.W_V(x)
        scores = q @ k.transpose(0, 1) / self.scale
        return torch.softmax(scores, dim=-1) @ v


torch.manual_seed(0)
model = TinyAttention(d_model=3).eval()

center = 0.2 * torch.randn(2, 3)
err = Error("input", ShapeDtype(center.shape, center.dtype))

for domain, interpret in [(Zono, zono_interpret), (PolySp, polysp_interpret)]:
    x = domain.error(err) * 0.02 + center
    lb, ub = interpret(model)(x).lbub()
    print(domain.__name__, "max width:", (ub - lb).max().item())

Notes

• The softmax handler decomposes \(\sigma_i = 1 / \sum_j e^{\nu_j - \nu_i}\) and supports the last axis (dim=-1).

• Products of two abstract operands go through the domain’s Matmul handler: an interval estimate for zonotopes, an exact second-order product (up to the configured truncation order) for sparse polynomials.

• Small perturbation radii are more stable for softmax-heavy models; validate soundness by sampling concrete inputs against lbub().