I wanted to see exactly how self-attention redistributes meaning across tokens without digging through megabytes of PyTorch C++ kernels. So I shipped a small Attention Inspector that implements scaled dot-product attention in NumPy, computes a real weight matrix over text embeddings from Oxlo.ai, and pipes the results to Llama 3.3 70B for a plain-language autopsy. If you are debugging long-context prompts or just want to touch the math that powers modern LLMs, this tool gets you there in under a hundred lines.
What you'll need
- Python 3.10+
pip install openai numpy- An Oxlo.ai API key from https://portal.oxlo.ai
Step 1: Set up the Oxlo.ai client
One client handles both the embedding and chat endpoints. I keep the key in an environment variable so I do not accidentally commit it.
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: Implement scaled dot-product attention
This is the exact operation inside every transformer block. I initialize random Q, K, V projections with a fixed seed so the demo stays reproducible.
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: Embed text and compute attention weights
I fetch BGE-Large embeddings through Oxlo.ai, project them down to 64 dimensions to keep the math small, then run the attention engine.
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: Define the inspector agent prompt
The system prompt tells the model to act like a debugging tool. I keep it strict so the output stays useful.
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: Send the pattern to Oxlo.ai for interpretation
I send the matrix to Llama 3.3 70B on Oxlo.ai. Because Oxlo.ai uses request-based pricing, it costs the same whether I send ten tokens or a thousand, so I can paste the whole matrix without counting context length.
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: Render an ASCII heatmap
Before reading the LLM's prose, I like to eyeball the raw scores. This helper prints a minimal heatmap in the terminal.
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)
Run it
Save everything in a file named attention_inspector.py, set your key, and run it.
export OXLO_API_KEY="sk-..."
python attention_inspector.py
The script first prints the heatmap, then the LLM analysis. Here is what the output looks like on my machine:
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.
Wrap-up
From here you can swap in Oxlo.ai's Qwen 3 32B to see if a different model family describes the same matrix differently, or you can replace the random projections with trained weights from a small open-source checkpoint. If you end up running this against long documents, keep in mind that Oxlo.ai's per-request pricing stays flat even when you pass thousands of tokens, which makes iterative prompt debugging far cheaper than token-based billing.
Top comments (0)