Model save

README.md

@@ -5,211 +5,48 @@ tags:
 model-index:
 - name: NeoLLM
   results: []
-license: apache-2.0
-datasets:
-- HuggingFaceFW/fineweb-edu
-language:
-- en
 ---
 
 <!-- This model card has been generated automatically according to the information the Trainer had access to. You
 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 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.
-
----
-
-## Architecture
-
-NeoLLM is a decoder-only transformer with the following configuration:
-
-| Parameter | Value |
-|---|---|
-| Hidden size | 512 |
-| Layers | 12 |
-| Attention heads | 8 |
-| KV heads (GQA) | 2 |
-| Head dim | 64 |
-| Intermediate size | 1536 |
-| Parameters | ~135M (~77M in tied embeddings, ~58M effective trainable) |
-| Vocabulary | Qwen3 tokenizer |
-| Context length | 512 tokens |
-
-### Integrated techniques
-
-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 and residual stream**
-
-- **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)) — 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.
-
----
-
-## Training
-
-| Setting | Value |
-|---|---|
-| Dataset | FineWeb-Edu |
-| 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% of steps) |
-| Weight decay | 0.1 |
-| Training time | ~6 hours |
-| Hardware | NVIDIA RTX 5090 (single GPU) |
-
-### Training curve
-
-| Step | Train Loss | Val Loss |
-|---|---|---|
-| 5,000 | 3.701 | 3.601 |
-| 10,000 | 3.286 | 3.169 |
-| 15,000 | 3.074 | 2.954 |
-| 20,000 | 2.935 | 2.808 |
-| 25,000 | 2.832 | 2.695 |
-| 30,000 | 2.747 | 2.605 |
-| 35,000 | 2.686 | 2.544 |
-| 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.
-
----
-
-## Benchmarks
-
-Evaluated with [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness), 0-shot.
-
-| Task | Metric | Score |
-|---|---|---|
-| MMLU | acc | 23.02% |
-| BoolQ | acc | 54.43% |
-| ARC-Easy | acc_norm | 34.55% |
-| ARC-Challenge | acc_norm | 22.44% |
-| PIQA | acc_norm | 53.59% |
-| HellaSwag | acc_norm | 26.78% |
-| OpenBookQA | acc_norm | 29.00% |
-| Winogrande | acc | 46.96% |
-| SCIQ | acc | 58.30% |
-| CommonsenseQA | acc | 20.88% |
-| Lambada (OpenAI) | acc | 9.02% |
-| TruthfulQA MC2 | acc | 46.45% |
-| COPA | acc | 57.00% |
-| ANLI R1 | acc | 35.90% |
-| ANLI R2 | acc | 33.90% |
-| ANLI R3 | acc | 33.58% |
-| Wikitext | word_ppl | 40.41 |
-| **Average (acc tasks)** | | **36.61%** |
 
----
-
-## 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.
-
-### Periodicity detection
-
-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.
-
-### Numerical stability
-
-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.
-
-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.
-
----
+# NeoLLM
 
-## Limitations
+This model is a fine-tuned version of [](https://huggingface.co/) on an unknown dataset.
 
-- **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.
+## Model description
 
----
+More information needed
 
-## Intended use
+## Intended uses & limitations
 
-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.
+More information needed
 
----
+## Training and evaluation data
 
-## References
-
-| Paper | arXiv |
-|---|---|
-| SeeDNorm: Self-Rescaled Dynamic Normalization | [2510.22777](https://arxiv.org/abs/2510.22777) |
-| Explicit Multi-head Attention (MEA) | [2601.19611](https://arxiv.org/abs/2601.19611) |
-| Learnable Multipliers | [2601.04890](https://arxiv.org/abs/2601.04890) |
-| Directional Routing in Transformers | [2603.14923](https://arxiv.org/abs/2603.14923) |
-| Exclusive Self Attention (XSA) | [2603.09078](https://arxiv.org/abs/2603.09078) |
-| Gated Attention for LLMs | [2505.06708](https://arxiv.org/abs/2505.06708) |
-| Affine-Scaled Attention | [2602.23057](https://arxiv.org/abs/2602.23057) |
-| The Curse of Depth in LLMs (LNS) | [2502.05795](https://arxiv.org/abs/2502.05795) |
-| LUCID: Attention with Preconditioned Representations | [2602.10410](https://arxiv.org/abs/2602.10410) |
-| FAN: Fourier Analysis Networks | [2502.21309](https://arxiv.org/abs/2502.21309) |
-| SimpleGPT | [2602.01212](https://arxiv.org/abs/2602.01212) |
-| 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) |
+More information needed
 
----
+## Training procedure
 
-## Author
+### Training hyperparameters
 
-[@Kyokopom](https://x.com/Kyokopom) on X
+The following hyperparameters were used during training:
+- learning_rate: 0.0006
+- train_batch_size: 64
+- eval_batch_size: 64
+- seed: 42
+- optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
+- lr_scheduler_type: linear
+- lr_scheduler_warmup_steps: 0.1
+- num_epochs: 1
 
----
+### Training results
 
-## License
 
-Apache 2.0
 
-## Citation
+### Framework versions
 
-```bibtex
-@misc{neollm2026,
-  title  = {NeoLLM: A Research Language Model Integrating Recent Attention and Normalization Techniques},
-  author = {KitsuVp},
-  year   = {2026},
-  url    = {https://huggingface.co/KitsuVp/NeoLLM}
-}
-```
+- Transformers 5.3.0
+- Pytorch 2.10.0+cu130
+- Datasets 4.8.4
+- Tokenizers 0.22.2

config.json

@@ -58,6 +58,7 @@
   "use_attn_res": false,
   "use_cache": false,
   "use_directional_routing": true,
+  "use_hadamard_o_proj": true,
   "use_jtokm": false,
   "use_lucid_attention": true,
   "use_mea_attention": true,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:202d68ac8bde26856d954faf7f52783058d0e5acded2ba4f68a66d1feab63fdb
-size 270564256
+oid sha256:37bc22e41799bea586823c8314b8cef63a91cba4e37249fac4892caa2e4c6d37
+size 264257256

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6c5dfcbd78dcb6c60f50339391222569a0781b6354aa89850658c03600129b49
+oid sha256:514630dbcee9ab34d950db77e5576aa3fceb753ae355487d0a5cb86076ee1701
 size 5329