We are building an attention visualizer that treats an LLM as an introspective attention layer. It reads a long document, takes a query, and returns the exact spans it focuses on, complete with synthetic attention scores. This is useful for debugging retrieval pipelines, teaching transformer internals, or building adaptive context windows.
What you'll need
- Python 3.10 or newer
pip install openai- An Oxlo.ai API key from https://portal.oxlo.ai
- A long text to analyze
Step 1: Prepare a long-context document and query
First we need a document long enough that naive keyword search would fail, plus a targeted query. I will use a six-paragraph technical summary of transformer architectures.
QUERY = "What is the time complexity of standard self-attention and how do sparse patterns help?"
DOCUMENT = """
The Transformer architecture, introduced in the seminal paper Attention Is All You Need, revolutionized natural language processing by replacing recurrent and convolutional layers with attention mechanisms. At its core lies the scaled dot-product attention, which computes a weighted sum of values based on the compatibility between queries and keys. This operation allows the model to draw global dependencies between input and output tokens regardless of their distance in the sequence. The computational complexity of this operation is O(n squared d) per layer, where n is the sequence length and d is the representation dimension. While powerful, this quadratic cost becomes a bottleneck when processing long documents or high-resolution inputs.
To address this limitation, researchers have proposed numerous efficient attention variants. Sparse attention patterns, such as those used in Longformer and BigBird, restrict each query to attend only to a fixed set of local and global tokens. Linear attention methods reformulate the softmax operation to achieve O(n d squared) or even O(n d) complexity by exploiting kernel feature maps and the associativity of matrix multiplication. More recently, state-space models and sub-quadratic architectures have emerged as alternatives that bypass attention entirely for very long sequences.
Multi-head attention extends the single attention function by projecting queries, keys, and values into multiple lower-dimensional subspaces. Each head can then attend to different representation axes, capturing syntactic, semantic, or positional relationships in parallel. The outputs of all heads are concatenated and linearly projected again, producing the final attention-augmented representation. This design is remarkably parallelizable and has proven essential for pretraining at scale.
In decoder-only autoregressive models, causal masking ensures that each position can only attend to itself and previous tokens. This prevents information leakage from future positions during training and inference. The attention scores are typically computed as the dot product of queries and keys, scaled by the square root of the key dimension, then passed through a softmax function to obtain normalized weights. Dropout is applied to these weights as a regularization measure before the weighted sum of values is calculated.
Modern large language models such as Llama, Qwen, and DeepSeek stack dozens of these attention layers, interleaved with feed-forward networks and normalization. The choice of positional encoding, whether absolute sinusoidal or rotary embeddings, significantly affects how attention distributes across sequence positions. Understanding these patterns is crucial for optimizing prompt engineering, context compression, and retrieval-augmented generation systems.
"""Step 2: Lock down the system prompt
We need the model to act like an attention head and emit structured data. The prompt below forces JSON output and prohibits paraphrasing so we can correlate scores back to the original text.
SYSTEM_PROMPT = """You are an attention visualization engine. Your job is to simulate the self-attention mechanism of a transformer.
Given a long document and a user query, identify the exact sentences in the document that are most relevant to answering the query. For each relevant sentence, assign an attention score from 0.0 to 1.0 representing how strongly a transformer would attend to that token span when answering the query.
Respond ONLY with a JSON object in this exact format:
{
"query": "",
"relevant_spans": [
{"span": "", "score": 0.95},
...
]
}
Rules:
- Do not modify the text of the spans. Copy them exactly from the document.
- Scores must be between 0.0 and 1.0.
- Include at most the top 5 most relevant spans.
- If no span is relevant, return an empty list.
"""Step 3: Wire the Oxlo.ai client
Now we wire the OpenAI SDK to Oxlo.ai's endpoint. I use Llama 3.3 70B because it handles the 128K context window needed for long documents. Because Oxlo.ai uses flat per-request pricing, I can pass the full document every time without the cost scaling with input length, which matters when you are iteratively debugging attention over long contexts. See https://oxlo.ai/pricing for details.
from openai import OpenAI
import json
client = OpenAI(base_url="https://api.oxlo.ai/v1", api_key="YOUR_OXLO_API_KEY")
def extract_attention(query, document):
user_message = f"Document:\n{document}\n\nQuery: {query}"
response = client.chat.completions.create(
model="llama-3.3-70b",
messages=[
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_message},
],
temperature=0.1,
max_tokens=1024,
)
raw = response.choices[0].message.content
# Strip markdown fences if the model wraps JSON in them
if raw.startswith("
```"):
raw = raw.split("\n", 1)[1].rsplit("```
", 1)[0]
return json.loads(raw.strip())Step 4: Render the attention heatmap
The model returns JSON. We map those exact spans back to the source lines and print an ASCII bar chart so we can see which sentences received the highest attention weights.
def render_attention_map(result, document):
spans = {item["span"]: item["score"] for item in result["relevant_spans"]}
lines = [line.strip() for line in document.splitlines() if line.strip()]
print(f"Query: {result['query']}\n")
print("Attention Map")
print("-" * 70)
for line in lines:
score = spans.get(line, 0.0)
bar = "█" * int(score * 20)
print(f"{score:.2f} | {bar:<20} | {line[:45]}")
attention_result = extract_attention(QUERY, DOCUMENT)
render_attention_map(attention_result, DOCUMENT)Run it
Running the script produces a clear map of where the model focuses its attention. The bars make it obvious which sentences are being weighted heavily for the answer.
Query: What is the time complexity of standard self-attention and how do sparse patterns help? Attention Map ---------------------------------------------------------------------- 0.00 | | The Transformer architecture, introduced in the seminal paper Attention Is All You Need, revolutionized natural language processing by replacing recurrent and convolutional layers with attention mechanisms. 0.95 | ████████████████████ | The computational complexity of this operation is O(n squared d) per layer, where n is the sequence length and d is the representation dimension. 0.92 | ███████████████████ | To address this limitation, researchers have proposed numerous efficient attention variants. 0.88 | ██████████████████ | Sparse attention patterns, such as those used in Longformer and BigBird, restrict each query to attend only to a fixed set of local and global tokens. 0.00 | | Multi-head attention extends the single attention function by projecting queries, keys, and values into multiple lower-dimensional subspaces. 0.00 | | In decoder-only autoregressive models, causal masking ensures that each position can only attend to itself and previous tokens. 0.00 | | Modern large language models such as Llama, Qwen, and DeepSeek stack dozens of these attention layers, interleaved with feed-forward networks and normalization.
Wrap-up
Swap in deepseek-v4-flash to stress-test attention over a 1M token context, or add a second Oxlo.ai call that asks the model to explain why each high-attention span is relevant. That turns the script into a lightweight RAG debugger you can drop into any pipeline.
Top comments (0)