Skip to content

Reverse Attention

Trace attention paths backward through transformer models, visualized as interactive Sankey diagrams.

  • Reverse Attention Tracing

    Discover which tokens most influence a target token by following attention paths backward through the model.

  • Beam Search

    Efficiently explore multiple high-probability paths through the attention matrix, not just the greedy maximum.

  • Interactive Visualization

    D3.js-powered Sankey diagrams with zoom, pan, click-to-highlight, and beam filtering.

  • HuggingFace Compatible

    Works with any HuggingFace transformer model that outputs attention weights.

Quick Example

from transformers import AutoModelForCausalLM, AutoTokenizer
from reverse_attention import ReverseAttentionTracer

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-0.5B",
    attn_implementation="eager",  # Required for attention output
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B")

# Create tracer and run
tracer = ReverseAttentionTracer(model, tokenizer)
result = tracer.trace_text("The quick brown fox jumps over the lazy dog.")

# Print attention paths
for path in result.paths_text:
    print(path)

# Generate interactive visualization
tracer.render_html(result, "output/", open_browser=True)

Why Reverse Attention?

When a transformer generates or processes text, each token "attends" to previous tokens. These attention weights tell us how much each token influenced the current one. But attention can chain—token 10 might attend to token 7, which attended heavily to token 3.

To truly understand what influenced your output, you need to trace these paths backward. That's exactly what this package does.

Use Cases

  • Debugging model behavior: Understand why a model generated a specific token
  • Interpretability research: Visualize attention flow patterns
  • Educational purposes: Learn how transformer attention works in practice

Installation

pip install reverse-attention

Or install from source:

git clone https://github.com/ovshake/rat
cd rat
pip install -e .

Next Steps