Update README.md

README.md

@@ -16,7 +16,7 @@ language:
 should probably proofread and complete it, then remove this comment. -->
 # NeoLLM
 
-NeoLLM is a 135M parameter language model trained from scratch on [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) in FP8. It integrates a collection of recently published attention and normalization techniques into a single architecture, with the goal of studying their combined behavior during pretraining. The model is actively being developed and the current checkpoint represents an intermediate training state.
+NeoLLM is a 135M parameter language model trained from scratch on [FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) in FP8 precision, completing training in approximately 6 hours on a single NVIDIA RTX 5090. It integrates a collection of recently published attention and normalization techniques into a single architecture, with the goal of studying how they interact during pretraining. The model is actively being developed and the current checkpoint represents an intermediate training state.
 
 ---
 
@@ -32,32 +32,36 @@ NeoLLM is a decoder-only transformer with the following configuration:
 | KV heads (GQA) | 2 |
 | Head dim | 64 |
 | Intermediate size | 1536 |
-| Parameters | ~135M (77M tied embeddings, ~58M effective) |
+| Parameters | ~135M (~77M in tied embeddings, ~58M effective trainable) |
 | Vocabulary | Qwen3 tokenizer |
-| Context length | 512 |
+| Context length | 512 tokens |
 
-### Components
+### Integrated techniques
 
-Each layer integrates the following mechanisms simultaneously:
+Each layer of NeoLLM combines the following mechanisms simultaneously. Rather than picking one improvement over another, the goal is to understand whether these techniques are compatible and complementary at small scale.
 
-**Normalization**
-- **SeeDNorm** ([arXiv:2510.22777](https://arxiv.org/abs/2510.22777)) — Self-rescaled dynamic normalization applied to Q and K projections
-- **PolyNorm / GPAS** ([arXiv:2506.22049](https://arxiv.org/abs/2506.22049), [arXiv:2602.04902](https://arxiv.org/abs/2602.04902)) — Polynomial composition activations with three branches (linear, quadratic, cubic) and gradient-preserving activation scaling on residual connections
-- **LayerNorm Scaling / LNS** ([arXiv:2502.05795](https://arxiv.org/abs/2502.05795)) — Per-layer scaling by 1/√ℓ to address the Curse of Depth
+**Normalization and residual stream**
 
-**Attention**
-- **FAN** ([arXiv:2502.21309](https://arxiv.org/abs/2502.21309)) — Fourier Analysis Networks: periodic channels in input projections for improved periodicity modeling
-- **MEA** ([arXiv:2601.19611](https://arxiv.org/abs/2601.19611)) — Explicit Multi-head Attention with learnable inter-head interaction matrices for K and V
-- **LUCID** ([arXiv:2602.10410](https://arxiv.org/abs/2602.10410)) — Lower-triangular preconditioner on V for attention decorrelation
-- **Affine-Scaled Attention** ([arXiv:2602.23057](https://arxiv.org/abs/2602.23057)) — Learnable per-head α and β scaling of softmax weights: `[α·softmax(QKᵀ)+β]V`
-- **XSA** ([arXiv:2603.09078](https://arxiv.org/abs/2603.09078)) — Exclusive Self Attention: removes the auto-position component from attention output
-- **Directional Routing** ([arXiv:2603.14923](https://arxiv.org/abs/2603.14923)) — Learned suppression of attention output along K=4 directions per head
-- **Gated Attention** ([arXiv:2505.06708](https://arxiv.org/abs/2505.06708)) — Sigmoid gate applied to attention output before o_proj
-- **Momentum Attention** ([arXiv:2411.03884](https://arxiv.org/abs/2411.03884)) — Causal first-difference shear on Q and K
+- **SeeDNorm** ([arXiv:2510.22777](https://arxiv.org/abs/2510.22777)) — Applied to Q and K projections. Dynamically rescales the normalization based on the input's own statistics, making the attention geometry more stable across varying input distributions.
+- **PolyNorm** ([arXiv:2602.04902](https://arxiv.org/abs/2602.04902)) — Replaces the standard MLP activation with three branches: linear (x), quadratic (x²), and cubic (x³) — each normalized and combined with learned weights. This allows the MLP to express both linear and non-linear relationships simultaneously.
+- **GPAS** ([arXiv:2506.22049](https://arxiv.org/abs/2506.22049)) — Gradient-Preserving Activation Scaling. Applied to the residual connections between sublayers; helps gradients flow more cleanly during training without distorting the residual stream.
+- **LayerNorm Scaling / LNS** ([arXiv:2502.05795](https://arxiv.org/abs/2502.05795)) — Each layer's output is scaled by 1/√ℓ where ℓ is the layer index. This directly addresses the "Curse of Depth" — the phenomenon where deeper layers in Pre-LN transformers contribute progressively less to the residual stream and effectively become redundant. With LNS, each layer maintains a meaningful contribution regardless of depth.
+
+**Attention mechanisms**
+
+- **FAN** ([arXiv:2502.21309](https://arxiv.org/abs/2502.21309)) — Fourier Analysis Networks. A portion of the input projection channels are dedicated to representing periodic patterns (cosine/sine pairs), while the remainder handle standard linear content. This helps the model detect recurring structure in sequences without relying solely on position embeddings.
+- **MEA** ([arXiv:2601.19611](https://arxiv.org/abs/2601.19611)) — Explicit Multi-head Attention. Adds small learnable interaction matrices between attention heads for K and V. In standard multi-head attention, each head operates independently; MEA allows heads to share information with each other before the attention computation.
+- **LUCID** ([arXiv:2602.10410](https://arxiv.org/abs/2602.10410)) — Applies a learned lower-triangular preconditioner to V (the values) before attention. This decorrelates the value representations across positions, reducing redundancy in what gets aggregated.
+- **Affine-Scaled Attention** ([arXiv:2602.23057](https://arxiv.org/abs/2602.23057)) — Adds two learnable per-head scalars (α and β) to the softmax weights: the effective attention becomes `[α·softmax(QKᵀ) + β]·V`. This allows the model to modulate how "peaked" vs "diffuse" each head's attention is, and to add a global context component via β·V that doesn't depend on the attention pattern.
+- **XSA** ([arXiv:2603.09078](https://arxiv.org/abs/2603.09078)) — Exclusive Self Attention. After computing attention, removes the component of the output that is aligned with the token's own value vector. This suppresses the tendency for tokens to "attend to themselves" and encourages each position to extract information from its context rather than its own representation.
+- **Directional Routing** ([arXiv:2603.14923](https://arxiv.org/abs/2603.14923)) — Each attention head learns K=4 directions in the output space. A learned router then decides how much to suppress the attention output along each direction, per input. This allows the model to selectively filter out interference that may appear in specific subspaces of the attention output.
+- **Gated Attention** ([arXiv:2505.06708](https://arxiv.org/abs/2505.06708)) — A sigmoid gate is applied to the attention output before the output projection. This introduces non-linearity and sparsity into the attention pathway, and helps prevent attention sinks (where one token absorbs most of the attention mass without contributing semantically).
+- **Momentum Attention** ([arXiv:2411.03884](https://arxiv.org/abs/2411.03884)) — Modifies Q and K by subtracting a fraction of the previous position's Q and K values (a causal first-difference). This makes the attention geometry sensitive to changes between consecutive positions rather than just their absolute values, acting as a local transition detector.
 
 **MLP**
-- **Learnable Multipliers** ([arXiv:2601.04890](https://arxiv.org/abs/2601.04890)) — Per-row and per-column learnable scale parameters on linear layers
-- **SimpleGPT** ([arXiv:2602.01212](https://arxiv.org/abs/2602.01212)) — Second-order geometry normalization applied inside MLP projections
+
+- **Learnable Multipliers** ([arXiv:2601.04890](https://arxiv.org/abs/2601.04890)) — Adds per-row and per-column learnable scalar parameters to each linear layer. Rather than having a fixed scale relationship between neurons, the model can freely adjust the relative importance of each row and column in every matrix independently of the weight values themselves.
+- **SimpleGPT** ([arXiv:2602.01212](https://arxiv.org/abs/2602.01212)) — A normalization strategy derived from second-order geometry analysis, applied inside the MLP projections to improve optimization stability.
 
 ---
 
@@ -66,13 +70,13 @@ Each layer integrates the following mechanisms simultaneously:
 | Setting | Value |
 |---|---|
 | Dataset | FineWeb-Edu |
-| Tokens | ~1.5B (45k steps × batch 64 × length 512) |
-| Precision | FP8 (E4M3/E5M2) native |
+| Tokens seen | ~1.5B (45k steps × batch 64 × length 512) |
+| Precision | FP8 native (E4M3 weights/activations, E5M2 gradients) |
 | Optimizer | AdamW fused (β₁=0.9, β₂=0.999, ε=1e-8) |
-| Learning rate | 6e-4 with linear warmup (10%) |
+| Learning rate | 6e-4 with linear warmup (10% of steps) |
 | Weight decay | 0.1 |
 | Training time | ~6 hours |
-| Hardware | NVIDIA RTX 5090 |
+| Hardware | NVIDIA RTX 5090 (single GPU) |
 
 ### Training curve
 
@@ -88,6 +92,9 @@ Each layer integrates the following mechanisms simultaneously:
 | 40,000 | 2.658 | 2.505 |
 | 45,000 | 2.493 | 2.345 |
 
+Val loss 2.345 corresponds to perplexity ≈ 10.4. Wikitext word-level perplexity at 45k steps: **40.4**.
+
+The training loss being consistently higher than validation loss indicates the model is in a healthy underfitting regime — it has not memorized the training data, and every additional token seen continues to be informative.
 
 ---
 
@@ -118,32 +125,49 @@ Evaluated with [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluati
 
 ---
 
-## Interpretability notes
+## Internal analysis
+
+A set of custom interpretability tools was developed alongside NeoLLM to measure how each architectural component behaves in practice. What follows is a summary of the most relevant findings from the current checkpoint, written to be readable without access to the tools themselves.
+
+### How the layers divide their work
+
+By measuring how much each layer changes the representation (angular distance between input and output hidden states), and which positions each attention head looks at, a clear functional hierarchy emerged across the 12 layers:
+
+- **Layers 0–3** perform aggressive local processing. Layer 0 makes the most radical transformation of any layer — the embedding representation exits almost orthogonal to where it entered. XSA is highly active here (removing self-referential components from attention output), which means these layers spend much of their energy cleaning up auto-position information rather than gathering context.
+- **Layers 4–6** detect patterns. Layer 5 in particular was found to attend to the preceding element in geometric sequences (e.g., in `1, 2, 4, 8`, position of `4` attends to `2`), suggesting genuine structural pattern recognition. The FAN periodic channels grow increasingly active through these layers when processing long, structured sequences.
+- **Layers 7–8** perform global integration. All 8 attention heads in these layers behave almost identically (cross-head similarity > 0.87), with the longest attention spans of the model. This is consistent with a consolidation phase where the model aggregates information from across the sequence before synthesis.
+- **Layers 9–11** prepare the final prediction. Layer 9 captures long-range correlations between tokens at distances of 13–22 positions. Layer 10 is the most active single layer in the model across nearly every metric simultaneously. Layer 11 selectively projects toward the vocabulary head, with its MLP contributing 81% of the final residual update.
+
+### Affine-Scaled Attention in practice
+
+The learned α values (which control how peaked or diffuse attention is) range from 0.23 to 0.54 across layers. A consistent pattern emerged: layers where XSA removes the most auto-position content tend to have lower α, effectively letting the global β·V term dominate. Layers where attention is more contextually meaningful have higher α. This suggests the two mechanisms — XSA and affine scaling — converged on a complementary division of labor that was not explicitly designed.
+
+Affine consistently increases attention entropy in every layer (ΔH > 0 in all cases), producing more distributed attention patterns. This is the intended behavior — reducing the model's tendency to focus exclusively on a single token — and is consistent with the benchmark improvement on tasks requiring broad context integration.
 
-A set of custom interpretability tools was developed alongside training to analyze the internal dynamics of the model. The following observations correspond to the 45k step checkpoint.
+### Periodicity detection
 
-**Layer hierarchy.** The model develops a functional hierarchy across layers. Early layers (L0–L3) focus on local structure processing, with XSA auto-position removal active at frac > 1.5 throughout. Middle layers (L4–L6) capture periodic and sequential patterns — L5 explicitly attends to preceding elements in geometric sequences. Layers L7–L8 perform global integration with cross-head similarity above 0.87. Layers L9–L11 perform synthesis toward the vocabulary head, with L11 MLP contributing rho_mlp = 0.81 to the residual stream.
+The FAN periodic channels become progressively more active with depth when processing long sequences, reaching R_FAN = 0.18 in the final layer (compared to ~0.08 in layer 0). With a short prompt (7 tokens), this gradient does not appear — the periodic channels activate in proportion to the amount of sequential structure available in the input. This is consistent with FAN's design intent.
 
-**Affine-Scaled Attention.** Learned α values range from 0.23 to 0.54 across layers. Lower α appears in layers with high XSA activity, where affine defers to global integration via the β·V_cumsum term. ΔH (attention entropy increase from affine) is consistently positive across all layers and inputs.
+### Numerical stability
 
-**Periodicity.** FAN R_FAN grows monotonically with depth in long sequences, reaching 0.18 in L11 at 130 tokens. This indicates that periodic channels are activated in proportion to available sequential structure in the input.
+One concern for FP8 training is the emergence of extremely large activation values (outliers), which can cause overflow. Using the colinearity analysis from the TWEO paper ([arXiv:2511.23225](https://arxiv.org/abs/2511.23225)) as a measurement framework, the MLP outlier risk score for NeoLLM is **1.96 on average**, compared to 500–5000 for standard transformers. The gate_proj and up_proj matrices maintain near-orthogonal dominant singular directions throughout training, which is why the model trains stably in FP8 without any additional engineering tricks.
 
-**Numerical condition.** MLP outlier risk (simulated via the colinearity mechanism described in TWEO, arXiv:2511.23225) is 1.96 mean (vs 500–5000 typical for standard transformers), reflecting near-orthogonal gate_proj and up_proj dominant directions. The only component with elevated condition numbers is o_proj (κ from 2,009 to 105,616), partially compensated by the learnable multipliers.
+The one component with elevated numerical condition is o_proj (the output projection of attention), where condition numbers range from ~2,000 to ~105,000 across layers. This is partially compensated by the learnable multipliers on those matrices, which learned to attenuate the high-energy directions. This will be explicitly regularized in future training runs.
 
 ---
 
 ## Limitations
 
-- The model is trained for a single epoch on FineWeb-Edu and has not reached its estimated token optimum. Knowledge-intensive tasks (MMLU, factual retrieval) remain limited.
-- A gradient spike near step 40k degraded the long-range correlation structure in layer 9, affecting tasks that require integrating information across long distances (Lambada, SCIQ, BoolQ). An earlier checkpoint around step 38k is expected to have better aggregate benchmark performance.
-- PolyNorm exclusivity for the quadratic branch (delta_12) is negative across all layers at 45k steps, indicating that the exclusive_logits parameters did not converge to effective orthogonalization. This will be addressed in subsequent runs by including exclusive_logits in the checkpoint from the start of training.
-- The model has not been instruction-tuned or aligned in any way. It is a base language model intended for research.
+- **Token budget.** The model has seen approximately 1.5B tokens, which is below its estimated optimum given its architecture. Performance on knowledge-intensive tasks (MMLU, factual recall) will continue to improve with more training.
+- **Gradient spike at step 40k.** A spike in gradient norm near step 40k reorganized the attention pattern in layer 9, which previously captured long-range token correlations. This degraded performance on tasks requiring long-range context integration (Lambada, SCIQ, BoolQ). A checkpoint from around step 38k is expected to have better aggregate benchmark scores.
+- **PolyNorm exclusivity.** The quadratic branch of PolyNorm was intended to learn representations orthogonal to the linear branch. At 45k steps, this orthogonality is not being maintained — the two branches have become partially redundant. This will be corrected in the next training run by including the relevant parameters in the checkpoint from the start.
+- **Base model only.** NeoLLM has not been instruction-tuned or aligned. It is a base language model trained purely for next-token prediction.
 
 ---
 
 ## Intended use
 
-NeoLLM is released for research purposes. It is intended as a testbed for studying the combined behavior of recent attention and normalization techniques during pretraining at small scale. It is not intended for production deployment in its current state.
+NeoLLM is released for research purposes. It is intended as a testbed for studying the combined behavior of recent attention and normalization techniques during pretraining at small scale. It is not intended for deployment in production systems in its current state. The training setup, interpretability tools, and architectural choices are actively being iterated on.
 
 ---
 
@@ -165,6 +189,7 @@ NeoLLM is released for research purposes. It is intended as a testbed for studyi
 | GPAS: Gradient-Preserving Activation Scaling | [2506.22049](https://arxiv.org/abs/2506.22049) |
 | PolyNorm / PolyCom | [2602.04902](https://arxiv.org/abs/2602.04902) |
 | Momentum Attention | [2411.03884](https://arxiv.org/abs/2411.03884) |
+| TWEO: Transformers Without Extreme Outliers (analysis reference) | [2511.23225](https://arxiv.org/abs/2511.23225) |
 
 ---