|
|
--- |
|
|
license: mit |
|
|
tags: |
|
|
- pytorch |
|
|
- addressed-state-attention |
|
|
- interpretable-ai |
|
|
- mechanistic-interpretability |
|
|
language: |
|
|
- en |
|
|
--- |
|
|
|
|
|
# Addressed State Attention (ASA) |
|
|
|
|
|
Interpretable slot-based attention achieving competitive language modeling performance. |
|
|
|
|
|
## Quick Start |
|
|
|
|
|
```python |
|
|
# Install directly from GitHub |
|
|
!pip install git+https://github.com/DigitalDaimyo/AddressedStateAttention.git |
|
|
|
|
|
from asa import load_asm_checkpoint, generate |
|
|
from transformers import AutoTokenizer |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
# Download checkpoint from Hugging Face |
|
|
ckpt_path = hf_hub_download( |
|
|
repo_id="DigitalDaimyo/AddressedStateAttention", |
|
|
filename="checkpoints/fineweb_187M_75k.pt" |
|
|
) |
|
|
|
|
|
# Load checkpoint |
|
|
model, cfg, ckpt = load_asm_checkpoint( |
|
|
ckpt_path, |
|
|
mode="analysis" |
|
|
) |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("gpt2") |
|
|
|
|
|
# Generate text |
|
|
print(generate(model, tokenizer, "Once upon a time")) |
|
|
|
|
|
Performance |
|
|
FineWeb, 187M params: 3.73 val loss / 41.6 PPL (75k steps•32 batch•1024 seq) |
|
|
Architecture: 21 layers, 768d, 12 heads, 16 slots |
|
|
Links |
|
|
Code: https://github.com/DigitalDaimyo/AddressedStateAttention |
|
|
Paper: https://github.com/DigitalDaimyo/AddressedStateAttention/paper_drafts |
