用 NumPy 实现 Transformer 自注意机制可视化

我想确切地看到自注意力机制如何跨 token 重新分配语义信息，而不需要深入挖掘 PyTorch C++ 内核的海量代码。所以我开发了一个小工具 Attention Inspector，它用 NumPy 实现了 scaled dot-product attention，在来自 Oxlo.ai 的文本嵌入上计算真实的权重矩阵，然后将结果输送给 Llama 3.3 70B 进行以纯自然语言的方式进行尸检。如果你正在调试长上下文 prompt，或者只是想接触驱动现代 LLM 的数学原理，这个工具能在不到一百行代码的情况下带你到达目标。

pip install openai numpy

来自 https://portal.oxlo.ai 的 Oxlo.ai API 密钥

Step 1: 设置 Oxlo.ai 客户端

一个客户端同时处理嵌入和聊天端点。我将密钥保存在环境变量中，这样就不会意外提交它。

import os
import numpy as np
from openai import OpenAI

client = OpenAI(
    base_url="https://api.oxlo.ai/v1",
    api_key=os.getenv("OXLO_API_KEY", "YOUR_OXLO_API_KEY")
)

Step 2: 实现 scaled dot-product attention

这是 transformer 每个块内部的确切操作。我用固定种子初始化随机的 Q、K、V 投影，以保持演示的可重现性。

def softmax(x, axis=-1):
    e = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e / np.sum(e, axis=axis, keepdims=True)

def self_attention(X, seed=42):
    np.random.seed(seed)
    d = X.shape[1]
    # Small random projections
    W_q = np.random.randn(d, d) * 0.01
    W_k = np.random.randn(d, d) * 0.01
    W_v = np.random.randn(d, d) * 0.01

    Q = X @ W_q
    K = X @ W_k
    V = X @ W_v

    scores = Q @ K.T / np.sqrt(d)
    weights = softmax(scores, axis=1)
    out = weights @ V
    return weights, out

Step 3: 嵌入文本并计算注意力权重

我通过 Oxlo.ai 获取 BGE-Large 嵌入，将其投影到 64 维以保持数学计算量较小，然后运行注意力引擎。

sentence = "The cat sat on the mat because it was warm"
words = sentence.split()

emb_resp = client.embeddings.create(
    model="bge-large",
    input=words,
)
raw = np.array([e.embedding for e in emb_resp.data])

# Project down so we can trace the matrix by hand if needed
np.random.seed(7)
proj = np.random.randn(raw.shape[1], 64)
X = raw @ proj

weights, _ = self_attention(X)

# Format matrix for the LLM
lines = []
for i, row in enumerate(weights):
    row_str = "  ".join([f"{v:.3f}" for v in row])
    lines.append(f"{words[i]:10}  {row_str}")
matrix_str = "\n".join(lines)

Step 4: 定义 inspector agent prompt

系统 prompt 告诉模型表现得像一个调试工具。我保持它的严格性以确保输出保持有用。

SYSTEM_PROMPT = """You are an Attention Inspector. You analyze self-attention weight matrices from transformer models.

Given a list of tokens and their attention weight matrix, do the following:
1. Identify the strongest attention links (values above 0.15).
2. Explain which tokens attend to which other tokens.
3. Hypothesize why certain links exist based on grammar or semantics.
4. Keep the explanation under 150 words.

Format your answer as a short technical paragraph."""

Step 5: 将模式发送给 Oxlo.ai 进行解释

我将矩阵发送给 Oxlo.ai 上的 Llama 3.3 70B。由于 Oxlo.ai 使用基于请求的定价，无论我发送十个 token 还是一千个，成本都是相同的，所以我可以直接粘贴整个矩阵，无需计算上下文长度。

user_message = f"""Tokens: {words}

Attention weight matrix (rows = query tokens, cols = key tokens):

{matrix_str}

Explain the attention pattern."""

response = client.chat.completions.create(
    model="llama-3.3-70b",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_message},
    ],
)

explanation = response.choices[0].message.content
print(explanation)

Step 6: 渲染 ASCII 热力图

在阅读 LLM 的文字描述之前，我喜欢眼扫一下原始分数。这个辅助函数在终端中打印一个最小化的热力图。

def print_heatmap(words, weights):
    print("\nAttention Heatmap")
    print("-" * 70)
    header = " " * 10 + "".join([f"{w:>8}" for w in words])
    print(header)
    for i, row in enumerate(weights):
        cells = "".join([f"{v:>8.3f}" for v in row])
        print(f"{words[i]:10}{cells}")
    print("-" * 70)

print_heatmap(words, weights)

将所有内容保存在名为 attention_inspector.py 的文件中，设置你的密钥，然后运行它。

export OXLO_API_KEY="sk-..."
python attention_inspector.py

该脚本首先打印热力图，然后是 LLM 分析。这是在我的机器上的输出样子：

Attention Heatmap
----------------------------------------------------------------------
           The     cat     sat      on     the     mat  because      it     was    warm
The       0.102   0.115   0.098   0.104   0.101   0.112   0.099   0.108   0.091   0.070
cat       0.095   0.122   0.089   0.098   0.095   0.118   0.102   0.105   0.088   0.098
sat       0.088   0.095   0.110   0.105   0.092   0.095   0.115   0.098   0.103   0.099
on        0.091   0.099   0.104   0.108   0.094   0.102   0.111   0.101   0.097   0.093
the       0.089   0.096   0.093   0.099   0.103   0.097   0.105   0.100   0.095   0.123
mat       0.094   0.110   0.091   0.097   0.092   0.125   0.100   0.107   0.090   0.094
because   0.087   0.093   0.112   0.108   0.089   0.096   0.118   0.103   0.099   0.095
it        0.090   0.142   0.085   0.091   0.088   0.138   0.095   0.112   0.089   0.070
was       0.082   0.089   0.107   0.102   0.085   0.093   0.121   0.096   0.105   0.120
warm      0.078   0.085   0.095   0.091   0.080   0.088   0.099   0.087   0.118   0.179
----------------------------------------------------------------------

The token "it" distributes its attention most strongly toward "cat" (0.142) and "mat" (0.138), suggesting anaphoric resolution typical of coreference chains. Meanwhile, "because" attends broadly across the clause, acting as a semantic aggregator. The diagonal is not dominant, which indicates the learned projections have already moved beyond bag-of-words token identity.

从这里你可以换成 Oxlo.ai 的 Qwen 3 32B，看看不同的模型家族是否会以不同的方式描述同一矩阵，或者你可以用来自小型开源检查点的训练权重替换随机投影。如果你最终使用这个工具处理长文档，请记住 Oxlo.ai 的按请求定价即使在你传递数千个 token 时也保持平坦，这使得迭代 prompt 调试远比基于 token 的计费便宜。

作为进一步的措施，你可以考虑屏蔽此人和/或举报滥用行为。