【村儿网通】把 Scaled Dot-Product Attention 展开写一遍

万恶之源

今天偶然之间看到这样一个公式

\[Attention(q_t, K, V) = \sum^m_{s=1} \frac{1}{Z}exp(\frac{<q_t,k_s>}{\sqrt d_k})\cdot v_s \]

突然满脑子都是问号，Attention的公式我记得：$$Attention(Q,K,V) = Softmax(\frac{Q\cdot K^\top}{\sqrt d_k})V$$
为啥会有个Z？这形式哪里来的？

只看一个query

一个老生常谈的比喻，Q代表用户提出的问题，K代表一个文档库，V代表和K一一对应的问题答案，Attention的计算过程，是比较用户问题和文档库，找到最相近的文章来回答用户问题。
在实际计算过程中，Q、K、V都是矩阵，我们假设$$Q\in R^{n\times d_k}, K\in R^{m\times d_k}, V\in R^{m\times d_v}$$
我们从Q中取出一条query，从K中取出一篇文章key：

\[q_t \in R^{d_k}， k_s \in R^{d_k} \]

那么这篇文章和query的“相关性”可以被表达为：

\[score(q_t, k_s) = \frac{<q_t. k_s^\top>}{\sqrt d_k} \]

有了“分数”还不够，分数可能是正的可能是负的，而我们需要的是能回答query的key，需要从K中找到“最重要的一个”，因此需要计算出每个key的“权重”：

\[a_{t,s} = Softmax_s(score(q_t, k_s))= \frac{exp(score(q_t,k_s))}{\sum^m_{j=1}exp(score(q_t, k_j))} \]

稍微注意一下K的维度，j要从1到m，然后我们可以计算出一篇文章的“结果”:

\[y_s = a_{t,s} \cdot v_s = \frac{exp(\frac{q_t \cdot k_s^\top}{\sqrt d_k})}{\sum_{j=1}^m exp(\frac{q_t\cdot k_j^\top}{\sqrt d_k})} \cdot v_s \]

那么全局的计算方式也就出来了：

\[y = \sum_{s=1}^m \frac{exp(\frac{q_t\cdot k_s^\top}{\sqrt d_k})}{\sum_{j=1}^m exp(\frac{q_t \cdot k_j^\top}{\sqrt d_k})} \cdot v_s \]

这时我们可以观察到分母中有一个大块头，甚至会发现这个大块头像是一个“常数”，如果我们换个元就会变成：

\[令Z= \sum_{j=1}^m exp(\frac{q_t\cdot k_j^\top}{\sqrt d_k}) \]

\[y = \sum_{s=1}^m \frac{1}{Z}\cdot exp(\frac{q_t\cdot k_s^\top}{\sqrt d_k})\cdot v_s \]

python实现计算过程

import numpy as np
import torch
from typing import Sequence, Union, Optional, Any

def make_tensor(
    shape: Sequence[int],
    mode: str = "rand",
    *,
    dtype: torch.dtype = torch.float32,
    device: Optional[Union[str, torch.device]] = None,
    low: float = 0.0,
    high: float = 1.0,
) -> torch.Tensor:
    """
    输入 shape（如 [1,2,3]），返回对应尺寸的张量。

    mode:
      - "rand":  [low, high) 均匀随机数（float）
      - "randn": 标准正态随机数（float）
      - "zeros": 全 0
      - "ones":  全 1
    """
    if isinstance(shape, torch.Size):
        shape = tuple(shape)
    else:
        shape = tuple(int(s) for s in shape)

    mode = mode.lower()
    if mode == "zeros":
        return torch.zeros(shape, dtype=dtype, device=device)
    if mode == "ones":
        return torch.ones(shape, dtype=dtype, device=device)
    if mode == "rand":
        return (high - low) * torch.rand(shape, dtype=dtype, device=device) + low
    if mode == "randn":
        return torch.randn(shape, dtype=dtype, device=device)

    raise ValueError(f"Unknown mode: {mode}. Use 'rand'/'randn'/'zeros'/'ones'.")

def print_matrix_info(x: Any, name: str = "x", max_items: int = 10) -> bool:
    """
    传入变量，判断是否为“矩阵/张量”(torch.Tensor 或 np.ndarray)。
    是则打印：类型、dtype、device(若有)、shape、维度、min/max/mean、部分元素预览。
    返回：是否为矩阵/张量。
    """
    is_torch = isinstance(x, torch.Tensor)
    is_numpy = isinstance(x, np.ndarray)

    if not (is_torch or is_numpy):
        print(f"{name}: not a tensor/ndarray (type={type(x)})")
        return False

    if is_torch:
        t = x.detach()
        info = {
            "type": type(x),
            "dtype": t.dtype,
            "device": t.device,
            "shape": tuple(t.shape),
            "ndim": t.ndim,
            "requires_grad": getattr(x, "requires_grad", False),
        }
        # 统计量（尽量避免对非浮点/空张量报错）
        if t.numel() > 0:
            tt = t.float() if not t.is_floating_point() else t
            info.update({
                "min": tt.min().item(),
                "max": tt.max().item(),
                "mean": tt.mean().item(),
            })
        else:
            info.update({"min": None, "max": None, "mean": None})

        print(f"{name}:")
        for k, v in info.items():
            print(f"  {k}: {v}")

        # 预览
        flat = t.flatten()
        n = min(flat.numel(), max_items)
        preview = flat[:n].cpu().tolist()
        print(f"  preview({n}/{flat.numel()}): {preview}")
        return True

    # numpy
    a = x
    info = {
        "type": type(x),
        "dtype": a.dtype,
        "shape": a.shape,
        "ndim": a.ndim,
    }
    if a.size > 0:
        aa = a.astype(np.float32, copy=False) if not np.issubdtype(a.dtype, np.floating) else a
        info.update({
            "min": float(aa.min()),
            "max": float(aa.max()),
            "mean": float(aa.mean()),
        })
    else:
        info.update({"min": None, "max": None, "mean": None})

    print(f"{name}:")
    for k, v in info.items():
        print(f"  {k}: {v}")

    flat = a.reshape(-1)
    n = min(flat.size, max_items)
    preview = flat[:n].tolist()
    print(f"  preview({n}/{flat.size}): {preview}")
    return True

N = 4 
M = 20
DIM_WORD = 8
DIM_V = 10
SEQUENCE_MAX_LEN = 15
BATCH_SIZE = 1 

matrix_Q = make_tensor([BATCH_SIZE, N, DIM_WORD], mode='rand')
# print_matrix_info(matrix_Q)
matrix_K = make_tensor([BATCH_SIZE, M, DIM_WORD], mode='rand')
# print_matrix_info(matrix_K)
matrix_V = make_tensor([BATCH_SIZE, M, DIM_V], mode='rand')
# print_matrix_info(matrix_V)

score = torch.matmul(matrix_Q, matrix_K.transpose(-2, -1))
score = score / (DIM_WORD ** 0.5)
print_matrix_info(score)
weight_q_k = torch.softmax(score, dim=-1)
print_matrix_info(weight_q_k)
answer = torch.matmul(weight_q_k, matrix_V)
print_matrix_info(answer)