Attention Mechanisms in LLMs

#learnai #oxlo #ai

We are going to build an Attention Inspector agent that reveals which parts of a long document an LLM uses to answer a question. This helps developers understand how attention mechanisms work in practice and why long-context performance matters.

What you'll need

Python 3.10 or newer
The openai SDK: pip install openai
An Oxlo.ai API key from https://portal.oxlo.ai

1. Connect to Oxlo.ai

Import the OpenAI SDK and point it at Oxlo.ai. I use kimi-k2.6 here because its 131K context window gives us room to demonstrate attention over a large input.

from openai import OpenAI

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

response = client.chat.completions.create(
    model="kimi-k2.6",
    messages=[{"role": "user", "content": "Hello"}],
)
print(response.choices[0].message.content)

2. Build a needle-in-a-haystack document

To inspect attention, we need a long text where critical facts are buried far from the question. I generate fifty sections of filler text and hide a single secret code in the middle.

def build_document():
    paragraphs = []
    for i in range(50):
        filler = f"Section {i}: Infrastructure review. " * 30
        paragraphs.append(filler.strip())
    # Hide the needle in the haystack
    paragraphs[25] += " The secret launch code for Project Mercury is 7294."
    return "\n\n".join(paragraphs)

long_doc = build_document()
print(f"Document length: {len(long_doc)} characters")

3. Write the Attention Inspector system prompt

The system prompt forces the model to return structured output showing exactly which spans it relied on, simulating an attention heatmap.

SYSTEM_PROMPT = """You are an Attention Inspector.
Answer the user's question using only the provided document.
After your answer, list the exact quotes from the document that you relied on, ordered by importance.
For each quote, give a relevance score between 0.0 and 1.0.

Return strictly JSON in this format:
{
  "answer": "string",
  "attention_spans": [
    {"quote": "string", "score": 0.99}
  ]
}
"""

4. Query the model with JSON mode

We stuff the full document into the user message, ask for the hidden code, and enable JSON mode so the response is machine readable.

import json

question = "What is the secret launch code for Project Mercury?"

user_message = f"""Document:
{long_doc}

Question: {question}"""

response = client.chat.completions.create(
    model="kimi-k2.6",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_message},
    ],
    response_format={"type": "json_object"},
    temperature=0.1,
)

result = json.loads(response.choices[0].message.content)
print(json.dumps(result, indent=2))

5. Render the attention map

Finally, we print each attended span with its score. High scores on the needle sentence prove the model successfully attended to the right location in the long context.

print(f"Answer: {result['answer']}\n")
print("Attention spans:")
for span in result["attention_spans"]:
    score = span["score"]
    quote = span["quote"].replace("\n", " ")[:120]
    print(f"  {score:.2f} | {quote}...")

Run it

Here is the complete script. Save it as attention_inspector.py, replace YOUR_OXLO_API_KEY, and run it.

from openai import OpenAI
import json

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

SYSTEM_PROMPT = """You are an Attention Inspector.
Answer the user's question using only the provided document.
After your answer, list the exact quotes from the document that you relied on, ordered by importance.
For each quote, give a relevance score between 0.0 and 1.0.

Return strictly JSON in this format:
{
  "answer": "string",
  "attention_spans": [
    {"quote": "string", "score": 0.99}
  ]
}
"""

def build_document():
    paragraphs = []
    for i in range(50):
        filler = f"Section {i}: Infrastructure review. " * 30
        paragraphs.append(filler.strip())
    paragraphs[25] += " The secret launch code for Project Mercury is 7294."
    return "\n\n".join(paragraphs)

long_doc = build_document()
question = "What is the secret launch code for Project Mercury?"

user_message = f"""Document:
{long_doc}

Question: {question}"""

response = client.chat.completions.create(
    model="kimi-k2.6",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_message},
    ],
    response_format={"type": "json_object"},
    temperature=0.1,
)

result = json.loads(response.choices[0].message.content)

print(f"Answer: {result['answer']}\n")
print("Attention spans:")
for span in result["attention_spans"]:
    score = span["score"]
    quote = span["quote"].replace("\n", " ")[:120]
    print(f"  {score:.2f} | {quote}...")

Expected output:

Answer: 7294

Attention spans:
  0.98 | Section 25: Infrastructure review. Infrastructure review. ... The secret launch code for Project Mercury is 7294....
  0.15 | Section 24: Infrastructure review. Infrastructure review. ....

Wrap-up

You now have a working agent that externalizes the attention process of a long-context LLM. Because Oxlo.ai charges a flat rate per request, running this experiment with 100K+ tokens costs the same as a one-sentence ping, which makes iterative tuning affordable.

Two concrete next steps: feed the agent a real codebase instead of synthetic text to see which functions are attended to during a refactoring question, or switch the model to deepseek-v3.2 to compare how different architectures allocate attention across the same context.