Update README.md

README.md

@@ -12,37 +12,19 @@ pipeline_tag: text-generation
 library_name: pytorch
 tags:
 - text-generation
-- text-generation-inference
 - small-language-model
 - slm
-- tiny-lm
 - from-scratch
-- pretrained-from-scratch
 - causal-lm
 - decoder-only
 - transformer
 - gqa
-- grouped-query-attention
 - rope
-- rotary-position-embeddings
 - swiglu
-- rmsnorm
-- tied-embeddings
-- muon-optimizer
-- wsd-schedule
-- warmup-stable-decay
+- exclusive-gated-attention
 - multi-token-prediction
-- mtp
 - triton
-- torch
 - pytorch
-- bf16
-- english
-- baseline
-- research
-- test
-- openweights
-- apache-2.0
 ---
 
 <div align="center">
@@ -89,6 +71,7 @@ tags:
 | Vocabulary | 16,384 BPE |
 | Layers | 16 |
 | Heads | 12 Q · 6 KV (GQA) |
+| Attention | **Exclusive Gated** |
 
 </td>
 <td valign="top" width="50%">
@@ -123,7 +106,7 @@ Pragya is a from-scratch small language model project. Every layer, kernel, and
 
 Concretely, this run was built to verify:
 
-- **Numerical correctness** — fused Triton kernels (attention post-processing, QK-norm + RoPE, add + RMSNorm, MoE group-GEMMs) match an eager reference to bf16 tolerance.
+- **Numerical correctness** — fused Triton kernels (exclusive gated attention post-processing, QK-norm + RoPE, add + RMSNorm, MoE group-GEMMs) match an eager reference to bf16 tolerance.
 - **Convergence stability** — no loss spikes, no diverging runs, no NaN gradients across 11,000 steps.
 - **Schedule sensitivity** — a longer Warmup–Stable–Decay (WSD) tail produces a sharp, measurable drop in loss.
 - **Data-budget sensitivity** — fewer unique tokens seen twice beats more unique tokens seen once, at fixed compute.
@@ -141,7 +124,9 @@ The model itself is disposable. The demonstration is not.
 | **Total parameters** | ~27M |
 | **Architecture** | Decoder-only transformer |
 | **Layers** | 16 |
-| **Attention** | 12 query heads · 6 key/value heads (GQA) |
+| **Attention** | Exclusive Gated Attention |
+| **Query heads** | 12 |
+| **Key/value heads** | 6 (GQA) |
 | **Hidden size** | 384 |
 | **Head dimension** | 32 |
 | **FFN** | SwiGLU (`ffn_mult = 8/3`) |
@@ -157,6 +142,49 @@ The model itself is disposable. The demonstration is not.
 
 ---
 
+## Attention: Exclusive Gated Attention
+
+The model uses **Exclusive Gated Attention** — a composition of two per-head gating operations, both always on, both zero-initialized to identity at step 0, applied to the SDPA output before the output projection.
+
+### What it is
+
+Standard multi-head attention computes `context = softmax(QKᵀ/√d)·V` and projects it. This model inserts two additional operations between the attention output and the output projection:
+
+**1. Gated Attention (Qiu et al., NeurIPS 2025)**
+
+A per-head sigmoid gate modulates the attention output. The gate reads a small window of the block's *input* (the first 16 hidden dims), producing one scalar per head:
+
+```
+gate  = 2 · sigmoid(W_gate · x[:16])
+context = context · gate
+```
+
+The `2 · sigmoid(·)` form means `gate = 1.0` at `W_gate = 0`, so a zero-initialized projection gives an exact identity op at step 0 — the head's output is passed through unmodified. Over training, each head learns to suppress (gate → 0) or amplify (gate → 2) its own contribution.
+
+**2. Exclusive Self-Attention (XSA)**
+
+After the gate, the model subtracts a scaled projection of the *self-value direction* from the context. This removes the component of the attention output that is just a copy of the value vector itself — "exclusive" because it excludes self-information:
+
+```
+v_unit     = v / ‖v‖
+proj_coeff = (context · v_unit)
+correction = tanh(α) · proj_coeff · v_unit
+context    = context − correction
+```
+
+`α` is a per-head learnable scalar, zero-initialized, so `tanh(α) = 0` at step 0 → identity. Each head learns how much (if any) self-value component to remove. **Position 0 is masked off** — at position 0, causal attention only sees itself, so subtracting the self component would remove 100% of the signal. The correction is zeroed at that position.
+
+### Why it matters at this scale
+
+Both gates are **zero-init to identity** — the model starts as vanilla attention and learns the gates incrementally. This avoids the training instability you'd get from random-initialized multiplicative gates at 21M parameters. The composite effect is that each head independently learns:
+
+- **How loud to be** (gate magnitude, 0 to 2×)
+- **How much of its output is redundant self-information** (XSA coefficient, −1 to +1)
+
+Both operations are fused into a single Triton kernel (`_attn_post_fwd_kernel` / `_attn_post_bwd_kernel`), so the effective cost over vanilla attention is roughly one extra elementwise pass over the context tensor, plus a small reduction.
+
+---
+
 ## Training
 
 <table>
@@ -226,8 +254,6 @@ Two checkpoints from the same training stack. Same model, same compute budget pe
 
 ### Difference table
 
-What changed between test2 and test3, at a glance:
-
 |  | test2 | test3 | Δ | Direction |
 |:---|:---:|:---:|:---:|:---:|
 | **Backbone (non-embedding)** | 21.25M | 21.25M | 0 | = identical |
@@ -332,7 +358,7 @@ At this size, tight sampling matters more than usual — the model has very litt
 ```
 pragya.pt                 ← checkpoint (state_dict + config)
 tokenizer/tokenizer.pkl   ← pickled tiktoken Encoding (vocab 16,384)
-model.py                  ← GPT architecture (GPT / GPTConfig)
+model.py                  ← GPT architecture (GPT / GPTConfig) + Exclusive Gated Attention
 inference.py              ← generate(), load_model(), sampling helpers
 ```
 
@@ -376,7 +402,7 @@ print(texts[0])
 
 `inference.py` handles checkpoint loading, KV-cached generation, the sampling stack (top-k / top-p / repetition penalty / no-repeat n-grams), and streaming output.
 
-`model.py` contains the full architecture — no external dependencies beyond PyTorch and the tokenizer.
+`model.py` contains the full architecture — including the fused Exclusive Gated Attention implementation — with no external dependencies beyond PyTorch and the tokenizer.
 
 ---
 
@@ -392,6 +418,12 @@ print(texts[0])
 
 ## FAQ
 
+<details>
+<summary><b>What is Exclusive Gated Attention?</b></summary>
+
+A composition of two per-head gates applied to the attention output: a **sigmoid gate** that scales each head's output (Qiu et al., NeurIPS 2025), and a **subtractive self-projection** that removes the component of the output that is just a copy of the value vector (Exclusive Self-Attention, XSA). Both gates are zero-initialized to identity, so the model starts as vanilla attention and learns the gates during training. See [Attention](#attention-exclusive-gated-attention) for the full explanation.
+</details>
+
 <details>
 <summary><b>Can I load this with <code>AutoModelForCausalLM.from_pretrained</code>?</b></summary>
 
@@ -430,7 +462,7 @@ If you use this checkpoint in research, please link to this repository. There is
 
 ```bibtex
 @misc{pragya-test3,
-  title  = {Pragya test3: a 27M from-scratch decoder-only language model},
+  title  = {Pragya test3: a 27M from-scratch decoder-only language model with Exclusive Gated Attention},
   author = {Arush Kumar},
   year   = {2026},
   note   = {Test checkpoint, not for production use},