OutrageouslyBad200/SMAT_ablations

SMAT — Semantic Attention

Trained checkpoints for SMAT (Semantic Attention), a transformer attention variant with a learnable semantic-similarity bias and per-token value gate.

• —Code: github.com/OutrageouslyBad200/SMATest
• —Architecture: 24 layers × 384d × 6 heads, block size 256, ~64 M parameters
• —Tokenizer: GPT-2 (tiktoken, vocab 50 257)
• —Training data: FineWeb-Edu sample-10BT, 98 M tokens
• —Training compute: 12 000 optimizer steps, batch 16 × grad_accum 2 (effective 32), RTX 4060

Equation

Attn(Q,K,V) = softmax(QK^T/sqrt(d_k) + λ·S + P + M) · (G ⊙ V)

• —S_ij = cos(W_s h_i, W_s h_j) — cosine similarity in shared projection
• —c_j = (1/n) Σ_{l≤j} S_jl — causal semantic centrality
• —G_j = σ(w_g^T h_j + μ·c_j + β) — per-token value gate
• —λ = softplus(λ_raw) — constrained positive scalar (per layer)

Repository contents

This HuggingFace repo hosts 20 checkpoints from the 5-seed ablation in Experiment 6 of the SMAT paper:

baseline_s0/final.pt   s_only_s0/final.pt   g_only_s0/final.pt   full_s0/final.pt
baseline_s1/final.pt   s_only_s1/final.pt   g_only_s1/final.pt   full_s1/final.pt
baseline_s2/final.pt   s_only_s2/final.pt   g_only_s2/final.pt   full_s2/final.pt
baseline_s3/final.pt   s_only_s3/final.pt   g_only_s3/final.pt   full_s3/final.pt
baseline_s4/final.pt   s_only_s4/final.pt   g_only_s4/final.pt   full_s4/final.pt

Each variant directory also contains config.json and metrics.jsonl (per-step training + eval logs). | Variant | use_S | use_G | Description | |---------|---------|---------|-------------| | baseline | False | False | Standard attention | | s_only | True | False | Semantic bias only | | g_only | False | True | Value gate only | | full | True | True | Full SMAT |

Results

Validation perplexity on FineWeb-Edu, 5 seeds, 12 000 steps: | Variant | Mean ppl | Std | Δ vs baseline | Seed wins | |---------|----------|-----|---------------|-----------| | Baseline | 79.75 | 1.69 | — | — | | S-only | 79.47 | 1.71 | −0.35% | 4/5 | | G-only | 79.02 | 1.65 | −0.90% | 5/5 | | Full SMAT | 78.65 | 1.75 | −1.37% | 5/5 | 0 NaN failures across 240 000 optimizer steps.

Usage

pip install torch numpy tiktoken huggingface_hub
git clone https://github.com/OutrageouslyBad200/SMATest.git
cd SMATest

Download a single checkpoint:

from huggingface_hub import hf_hub_download
ckpt_path = hf_hub_download(
    repo_id="OutrageouslyBad200/SMATest",
    filename="full_s0/final.pt",
)

Load it into the SMAT model:

import torch
from model import Config, SMATTransformer
state = torch.load(ckpt_path, map_location="cuda")
cfg = Config(**state["config"])
model = SMATTransformer(cfg).cuda()
model.load_state_dict(state["state_dict"])
model.eval()

Reproduce surgical ablations (Experiment 7):

python ablate.py --ckpt full_s0/final.pt --n_batches 80

Surgical-ablation findings (Experiment 7)

Run on Full SMAT, val ppl 79.010:

• —The gate G is catastrophically essential.
• —S routes mostly through μ·c in the gate (74 % of lift), not through λ·S in attention (26 %).
• —Per-token gate differentiation matters: replacing G with its mean costs 149 %.

Limitations

• —Small base model (~64 M params); larger-scale runs (100 M on FineWeb / FineMath) show stronger lifts (−11 % to −17 %) but are not included as released checkpoints.
• —Trained only on English FineWeb-Edu sample-10BT — generalization to other domains untested at this scale.
• —Not instruction-tuned, not RLHF'd, no safety filtering. Research artifact only.

Citation

@misc{smat2026,
  author       = {OutrageouslyBad200},
  title        = {SMAT: Semantic Attention},
  year         = {2026},
  howpublished = {\url{https://github.com/OutrageouslyBad200/SMATest}},
}

Contact

For further information on training runs, intermediate experiments, or the unpublished paper draft, please contact the creator via GitHub or HuggingFace.

License

MIT License.