FiveAttn: 5-Component Attention & MHA 100M Transformers

A production-grade PyTorch implementation and benchmark suite comparing 5-Component Attention (Gated Attention Unit / GAU) against standard Multi-Head Attention (MHA) at the 100M-parameter scale, built for high-performance training on the TinyStories dataset.

---

Highlights

• 5-Component Attention (GAU): Replaces standard 4-projection MHA ($W_q, W_k, W_v, W_o$) with 5 projection matrices ($W_q, W_k, W_v, W_g, W_o$), introducing a non-linear Swish gate ($G = X W_g$) that modulates the attention context prior to output projection.
• Strict 100M Parameter Parity: Both the standard MHA and 5-Component models are calibrated to exactly 99,994,176 parameters (within 0.006% of 100M), guaranteeing fair scientific ablation.
• SmolLM-135M Tokenizer: Uses Hugging Face's HuggingFaceTB/SmolLM-135M tokenizer with a 49,152 vocabulary size and tied word embeddings.
• Apple Cut-Cross-Entropy: Computes memory-efficient next-token cross-entropy loss directly from final hidden states without allocating the $(B \times T \times V)$ logits matrix into memory, saving gigabytes of VRAM.
• SwiGLU & RMSNorm: Uses Swish-Gated Linear Units (SwiGLU) for the feed-forward network and Root Mean Square Normalization (RMSNorm) with pre-normalization residuals.
• Zero Graph Breaks in torch.compile: Verified with torch._dynamo.explain to compile into 1 single graph with 0 graph breaks.
• PyTorch Inductor Autotune Ready: Supports --compile --autotune (mode="max-autotune" and coordinate descent tuning).
• Remote-First CLI: Download and pre-tokenize TinyStories on your remote GPU instance using memory-mapped binary files (.bin) for high-throughput zero-copy training.

---

Architecture & Math

1. Multi-Head Attention (Standard 4 Projections)

Standard MHA projects input $X \in \mathbb{R}^{B \times T \times D}$ into queries, keys, and values: $$Q = X W_q, \quad K = X W_k, \quad V = X W_v$$ $$A = \text{softmax}\left(\frac{Q K^T}{\sqrt{d_k}} + \text{causal_mask}\right) V$$ $$\text{Output} = A W_o$$ Parameters per attention block: $4 \times D^2$.

2. 5-Component Attention (Gated Attention Unit - GAU)

5-Component Attention introduces a 5th projection matrix $W_g \in \mathbb{R}^{D \times D}$ acting as a data-dependent gating channel: $$Q = X W_q, \quad K = X W_k, \quad V = X W_v, \quad G = X W_g$$ $$A = \text{softmax}\left(\frac{\text{RoPE}(Q) \text{RoPE}(K)^T}{\sqrt{d_k}} + \text{causal_mask}\right) V$$ $$Z = A \odot \text{SiLU}(G)$$ $$\text{Output} = Z W_o$$ Just as SwiGLU adds a gating projection to standard MLPs, 5-Component Attention adds a gating projection to the attention mechanism. Parameters per attention block: $5 \times D^2$.

3. SwiGLU MLP

$$\text{SwiGLU}(X) = \left(\text{SiLU}(X W_{\text{gate}}) \odot (X W_{\text{up}})\right) W_{\text{down}}$$

4. RMSNorm

$$\text{RMSNorm}(X) = \frac{X}{\sqrt{\frac{1}{D}\sum_{i=1}^D X_i^2 + \epsilon}} \odot \gamma$$

---

100M Parameter Calibration Table

Both models share identical embedding size, depth (16 layers), attention heads (9 heads, 64 dim), and total parameters:

Component	Standard MHA-100M	5-Component Attention 100M	Difference
Vocab Size	49,152 (SmolLM)	49,152 (SmolLM)	0
Layers ($L$)	16	16	0
Hidden Dimension ($D$)	576	576	0
Attention Heads ($H$)	9	9	0
Head Dimension ($D_h$)	64	64	0
Attention Projections	4 ($W_q, W_k, W_v, W_o$)	5 ($W_q, W_k, W_v, W_g, W_o$)	+1 matrix
MLP Hidden Dimension	1,824	1,632	-192
Embedding Parameters	28,311,552	28,311,552	0
Attention Parameters	21,233,664	26,542,080	+5,308,416
SwiGLU MLP Parameters	50,429,952	45,121,536	-5,308,416
RMSNorm Parameters	19,008	19,008	0
Total Parameters	99,994,176	99,994,176	0 (EXACT MATCH)

---

Installation

# Clone repository
git clone https://github.com/your-org/fiveattn.git
cd fiveattn

# Install in editable mode
pip install -e .

# Or install with dev dependencies
pip install -e ".[dev]"

---

Remote Workflow (Recommended)

To run your training on a remote GPU machine (e.g. Lambda, RunPod, AWS, GCP):

Step 1: Pre-tokenize TinyStories on Remote

Run prepare-data on your remote server to download and convert TinyStories into high-throughput memory-mapped binary files:

# Tokenize train split (~2.1M stories, uint16 memmap)
fiveattn prepare-data --output-dir data/tinystories --split train

# Tokenize validation split
fiveattn prepare-data --output-dir data/tinystories --split validation

Step 2: Train 5-Component Attention (100M)

fiveattn train \
  --model fiveattn \
  --data-path data/tinystories/train.bin \
  --val-data-path data/tinystories/validation.bin \
  --batch-size 32 \
  --seq-len 512 \
  --lr 5e-4 \
  --steps 5000 \
  --device cuda \
  --compile \
  --autotune

Step 3: Train Standard MHA Baseline (100M)

fiveattn train \
  --model mha \
  --data-path data/tinystories/train.bin \
  --val-data-path data/tinystories/validation.bin \
  --batch-size 32 \
  --seq-len 512 \
  --lr 5e-4 \
  --steps 5000 \
  --device cuda \
  --compile \
  --autotune

---

CLI Command Reference

fiveattn count-params

Prints detailed parameter breakdown for both models:

fiveattn count-params

fiveattn verify-compile

Formally checks with torch._dynamo.explain that models have 0 graph breaks:

fiveattn verify-compile --seq-len 64

fiveattn benchmark

Runs latency, peak VRAM, and tokens/sec throughput benchmarks comparing MHA vs 5-Component Attention:

fiveattn benchmark --batch-size 8 --seq-len 256 --steps 50 --device cuda --compile

fiveattn train

Trains a model on TinyStories (online, offline memmap, or synthetic):

# Quick offline smoke-test on CPU with synthetic data
fiveattn train --model fiveattn --steps 10 --batch-size 2 --seq-len 64 --device cpu --synthetic

# Full training on CUDA with compiled kernels and autotuning
fiveattn train --model fiveattn --steps 2000 --batch-size 16 --seq-len 512 --device cuda --compile --autotune

fiveattn generate

Generates text continuation from a checkpoint or initialized model:

fiveattn generate \
  --checkpoint checkpoints/fiveattn_step_2000.pt \
  --prompt "Once upon a time, there was a little robot" \
  --max-tokens 100 \
  --temperature 0.8

---

Python API Usage

import torch
from fiveattn import ModelConfig, TransformerLM, CutCrossEntropyLoss

# 1. Initialize 100M 5-component attention model
config = ModelConfig.get_100m_fiveattn_config()
model = TransformerLM(config).cuda()

# 2. Forward pass with cut-cross-entropy loss
loss_fn = CutCrossEntropyLoss(shift=True, impl="cce")

input_ids = torch.randint(0, config.vocab_size, (4, 512), device="cuda")
targets = input_ids.clone()

# Directly compute hidden states and linear cross-entropy without allocating (B, T, V) logits
hidden_states = model.forward_hidden_states(input_ids)
loss = loss_fn(hidden_states, model.get_output_embeddings(), targets)
loss.backward()

# 3. Autoregressive generation
from fiveattn import generate_text
output = generate_text(model, prompt="Once upon a time")
print(output)

---

Running Unit Tests

pytest -v

All 15 tests run fully offline without network access or dataset downloads.

---

License

Apache-2.0