Model save

README.md

@@ -1,203 +1,60 @@
 ---
-language: en
-license: apache-2.0
+library_name: transformers
 tags:
-- causal-lm
-- research
-- fp8
-- attention
-- normalization
-- neollm
-datasets:
-- HuggingFaceFW/fineweb-edu
+- generated_from_trainer
+model-index:
+- name: NeoLLM
+  results: []
 ---
 
-# NeoLLM
-
-NeoLLM is a **135 M parameter** decoder-only 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.
+<!-- 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. -->
 
-> **Author / contact:** [@Kyokopom](https://x.com/Kyokopom) on X
-> **Repository:** [KitsuVp/NeoLLM](https://huggingface.co/KitsuVp/NeoLLM)
-
----
-
-## 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 |
-| Vocabulary | Qwen3 tokenizer (200,005 tokens) |
-| Context length | 512 tokens |
-
-### Parameter breakdown
-
-| Parameter bucket | Count |
-|---|---|
-| **Total parameters** | 158.16M (158,156,792) |
-| **Embedding parameters** (tied) | 102.40M (102,402,560) |
-| **Non-embedding parameters** | 55.75M (55,754,232) |
-| **Effective trainable parameters** | 158.16M (158,156,792) |
-
-> Weight tying is **enabled**: the input embedding matrix and the language-model head
-> share the same parameters, so the effective trainable budget is
-> `total − embed = 55.75M`.
-
-### Integrated techniques
-
-Each layer combines the following mechanisms simultaneously.
-
-**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 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. Directly addresses the
-  "Curse of Depth" in Pre-LN transformers.
-
-**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.
-- **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.
-- **LUCID** ([arXiv:2602.10410](https://arxiv.org/abs/2602.10410)) — Applies a learned
-  lower-triangular preconditioner to V before attention, decorrelating value representations
-  across positions.
-- **Affine-Scaled Attention** ([arXiv:2602.23057](https://arxiv.org/abs/2602.23057)) — Adds
-  two learnable per-head scalars (α and β) to the softmax weights:
-  `[α·softmax(QKᵀ) + β]·V`.
-- **XSA** ([arXiv:2603.09078](https://arxiv.org/abs/2603.09078)) — Exclusive Self Attention.
-  After computing attention, removes the component of the output aligned with the token's
-  own value vector.
-- **Directional Routing** ([arXiv:2603.14923](https://arxiv.org/abs/2603.14923)) — Each head
-  learns K=4 directions in the output space; a learned router suppresses the attention output
-  along each direction per input.
-- **Gated Attention** ([arXiv:2505.06708](https://arxiv.org/abs/2505.06708)) — A sigmoid gate
-  is applied to the attention output before the output projection, introducing non-linearity
-  and preventing attention sinks.
-- **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 (causal
-  first-difference).
-
-**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.
-- **SimpleGPT** ([arXiv:2602.01212](https://arxiv.org/abs/2602.01212)) — A normalization
-  strategy derived from second-order geometry analysis, applied inside MLP projections to
-  improve optimization stability.
-
----
+# NeoLLM
 
-## Training
-
-| Setting | Value |
-|---|---|
-| Dataset | FineWeb-Edu (sample-10BT) |
-| Tokens seen | ~1.54B (46,875 steps × batch 64 × length 512) |
-| Precision | FP8 native (E4M3 weights/activations, E5M2 gradients) + BF16 fallback |
-| Optimizer | Conda (Column-Normalized Adam) + GPA |
-| Learning rate | 6e-04 with linear warmup (10 % of steps) |
-| Weight decay | 0.1 |
-| Training time | ~6h 11m |
-| Hardware | NVIDIA RTX 5090 (single GPU) |
-
-### Training curve
-
-| Step | Train Loss | Val Loss |
-|---|---|---|
-| 5,000 | 6.297 | 6.164 |
-| 10,000 | 5.631 | 5.491 |
-| 15,000 | 5.333 | 5.188 |
-| 20,000 | 5.158 | 4.997 |
-| 25,000 | 5.033 | 4.863 |
-| 30,000 | 4.933 | 4.761 |
-| 35,000 | 4.861 | 4.685 |
-| 40,000 | 4.761 | 4.598 |
-| 45,000 | 4.698 | 4.524 |
-| 46,875 | — | 4.507 |
+This model is a fine-tuned version of [](https://huggingface.co/) on an unknown dataset.
+It achieves the following results on the evaluation set:
+- Loss: 3.6334
 
----
+## Model description
 
-## Limitations
+More information needed
 
-- **Token budget** — ~1.5 B tokens seen; below estimated optimum. Knowledge-intensive tasks
-  will improve with more training.
-- **Gradient spike at step 40k** — Reorganized the attention pattern in layer 9 that
-  previously captured long-range token correlations. A checkpoint from ~step 38k is expected
-  to have better aggregate benchmark scores.
-- **PolyNorm exclusivity** — The quadratic branch has become partially redundant with the
-  linear branch. Will be corrected in the next training run.
-- **Base model only** — Not instruction-tuned or aligned; purely a next-token-prediction
-  base model.
+## Intended uses & limitations
 
----
+More information needed
 
-## References
-
-All papers whose techniques are integrated into NeoLLM's architecture:
-
-| Technique | Paper title | arXiv |
-|---|---|---|
-| SeeDNorm | Self-Rescaled Dynamic Normalization | [2510.22777](https://arxiv.org/abs/2510.22777) |
-| MEA | Explicit Multi-head Attention | [2601.19611](https://arxiv.org/abs/2601.19611) |
-| Learnable Multipliers | Freeing the Scale of Language Model Matrix Layers | [2601.04890](https://arxiv.org/abs/2601.04890) |
-| Directional Routing | Directional Routing in Transformers | [2603.14923](https://arxiv.org/abs/2603.14923) |
-| XSA | Exclusive Self Attention | [2603.09078](https://arxiv.org/abs/2603.09078) |
-| Gated Attention | Gated Attention for LLMs | [2505.06708](https://arxiv.org/abs/2505.06708) |
-| Affine-Scaled Attention | Affine-Scaled Attention | [2602.23057](https://arxiv.org/abs/2602.23057) |
-| LNS | The Curse of Depth in LLMs | [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 | SimpleGPT | [2602.01212](https://arxiv.org/abs/2602.01212) |
-| GPAS | Gradient-Preserving Activation Scaling | [2506.22049](https://arxiv.org/abs/2506.22049) |
-| PolyNorm | PolyNorm / PolyCom | [2602.04902](https://arxiv.org/abs/2602.04902) |
-| Momentum Attention | Momentum Attention | [2411.03884](https://arxiv.org/abs/2411.03884) |
-| TWEO (analysis ref.) | Transformers Without Extreme Outliers | [2511.23225](https://arxiv.org/abs/2511.23225) |
+## Training and evaluation data
 
----
+More information needed
 
-## Citation
+## Training procedure
 
-```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}
-}
-```
+### Training hyperparameters
 
----
+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
 
-## Author
+### Training results
 
-[@Kyokopom](https://x.com/Kyokopom) on X
+| Training Loss | Epoch | Step  | Validation Loss |
+|:-------------:|:-----:|:-----:|:---------------:|
+| 4.0003        | 0.32  | 5000  | 3.9566          |
+| 3.7787        | 0.64  | 10000 | 3.7377          |
+| 3.6911        | 0.96  | 15000 | 3.6386          |
+| 3.6809        | 1.0   | 15625 | 3.6334          |
 
----
 
-## License
+### Framework versions
 
-Apache 2.0
+- Transformers 5.5.3
+- Pytorch 2.11.0+cu130
+- Datasets 4.8.4
+- Tokenizers 0.22.2

chat_template.jinja

@@ -1,54 +1,45 @@
-{%- if tools %}
-    {{- '<|im_start|>system\n' }}
-    {%- if messages[0]['role'] == 'system' %}
-        {{- messages[0]['content'] }}
-    {%- else %}
-        {{- 'You are a helpful assistant.' }}
-    {%- endif %}
-    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
-    {%- for tool in tools %}
-        {{- "\n" }}
-        {{- tool | tojson }}
-    {%- endfor %}
-    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
-{%- else %}
-    {%- if messages[0]['role'] == 'system' %}
-        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
-    {%- else %}
-        {{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
-    {%- endif %}
-{%- endif %}
-{%- for message in messages %}
-    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
-        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
-    {%- elif message.role == "assistant" %}
-        {{- '<|im_start|>' + message.role }}
-        {%- if message.content %}
-            {{- '\n' + message.content }}
-        {%- endif %}
-        {%- for tool_call in message.tool_calls %}
-            {%- if tool_call.function is defined %}
-                {%- set tool_call = tool_call.function %}
-            {%- endif %}
-            {{- '\n<tool_call>\n{"name": "' }}
-            {{- tool_call.name }}
-            {{- '", "arguments": ' }}
-            {{- tool_call.arguments | tojson }}
-            {{- '}\n</tool_call>' }}
-        {%- endfor %}
-        {{- '<|im_end|>\n' }}
-    {%- elif message.role == "tool" %}
-        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
-            {{- '<|im_start|>user' }}
-        {%- endif %}
-        {{- '\n<tool_response>\n' }}
-        {{- message.content }}
-        {{- '\n</tool_response>' }}
-        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
-            {{- '<|im_end|>\n' }}
-        {%- endif %}
-    {%- endif %}
-{%- endfor %}
-{%- if add_generation_prompt %}
-    {{- '<|im_start|>assistant\n' }}
-{%- endif %}
+{{- bos_token -}}
+{%- set keep_past_thinking = keep_past_thinking | default(false) -%}
+{%- set ns = namespace(system_prompt="") -%}
+{%- if messages[0]["role"] == "system" -%}
+    {%- set ns.system_prompt = messages[0]["content"] -%}
+    {%- set messages = messages[1:] -%}
+{%- endif -%}
+{%- if tools -%}
+    {%- set ns.system_prompt = ns.system_prompt + ("\n" if ns.system_prompt else "") + "List of tools: [" -%}
+    {%- for tool in tools -%}
+        {%- if tool is not string -%}
+            {%- set tool = tool | tojson -%}
+        {%- endif -%}
+        {%- set ns.system_prompt = ns.system_prompt + tool -%}
+        {%- if not loop.last -%}
+            {%- set ns.system_prompt = ns.system_prompt + ", " -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {%- set ns.system_prompt = ns.system_prompt + "]" -%}
+{%- endif -%}
+{%- if ns.system_prompt -%}
+    {{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
+{%- endif -%}
+{%- set ns.last_assistant_index = -1 -%}
+{%- for message in messages -%}
+    {%- if message["role"] == "assistant" -%}
+        {%- set ns.last_assistant_index = loop.index0 -%}
+    {%- endif -%}
+{%- endfor -%}
+{%- for message in messages -%}
+    {{- "<|im_start|>" + message["role"] + "\n" -}}
+    {%- set content = message["content"] -%}
+    {%- if content is not string -%}
+        {%- set content = content | tojson -%}
+    {%- endif -%}
+    {%- if message["role"] == "assistant" and not keep_past_thinking and loop.index0 != ns.last_assistant_index -%}
+        {%- if "</think>" in content -%}
+            {%- set content = content.split("</think>")[-1] | trim -%}
+        {%- endif -%}
+    {%- endif -%}
+    {{- content + "<|im_end|>\n" -}}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+    {{- "<|im_start|>assistant\n" -}}
+{%- endif -%}

config.json

@@ -11,12 +11,12 @@
     "AutoModel": "modeling_neollm.NeoLLMModel",
     "AutoModelForCausalLM": "modeling_neollm.NeoLLMForCausalLM"
   },
-  "bos_token_id": 200004,
+  "bos_token_id": 1,
   "directional_routing_k": 4,
   "directional_routing_temp": 3.0,
   "dropout_rate": 0.1,
   "dtype": "bfloat16",
-  "eos_token_id": 200004,
+  "eos_token_id": 7,
   "fan_ratio": 0.125,
   "fan_ratio_ffn": 0.0625,
   "generator_d_seed": 128,
@@ -45,7 +45,7 @@
   "num_attention_heads": 8,
   "num_hidden_layers": 12,
   "num_key_value_heads": 2,
-  "pad_token_id": 200001,
+  "pad_token_id": 0,
   "partial_rotary_factor": 0.25,
   "polynorm_exclusive": false,
   "repo_d_p": 64,
@@ -58,16 +58,16 @@
   },
   "rope_theta": 10000.0,
   "tie_word_embeddings": true,
-  "transformers_version": "5.5.0",
+  "transformers_version": "5.5.3",
   "use_affine_scaled_attention": true,
   "use_attn_res": false,
   "use_cache": false,
-  "use_directional_routing": true,
+  "use_directional_routing": false,
   "use_hadamard_o_proj": true,
   "use_jtokm": false,
-  "use_laurel": true,
-  "use_laurel_lr": true,
-  "use_laurel_rw": true,
+  "use_laurel": false,
+  "use_laurel_lr": false,
+  "use_laurel_rw": false,
   "use_lucid_attention": true,
   "use_mea_attention": true,
   "use_momentum_attention": true,
@@ -83,6 +83,6 @@
   "versatile_gumbel_temp_start": 5.0,
   "versatile_max_depth": 2,
   "versatile_total_experts": 4,
-  "vocab_size": 200005,
+  "vocab_size": 64402,
   "xsa_eps": 1e-06
 }

generation_config.json

@@ -1,11 +1,11 @@
 {
   "_from_model_config": true,
-  "bos_token_id": 200004,
+  "bos_token_id": 1,
   "eos_token_id": [
-    200004
+    7
   ],
   "output_attentions": false,
   "output_hidden_states": false,
-  "pad_token_id": 200001,
-  "transformers_version": "5.5.0"
+  "pad_token_id": 0,
+  "transformers_version": "5.5.3"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4f99fdeed60c04f6ed6ac81513606c7c2a97a2d7d30c6e98f4e22716050066ed
-size 316396360
+oid sha256:ca5fb5b2d4172761f6f863a1d57b1fab971b8c411b3239bd2d8925fb4bf171ea
+size 156549256

tokenizer.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4ad7e65f213ca8139ec0ef1f3445de27c21a6fecdad89d50b2af23efd6bd7d76
-size 14937008
+oid sha256:df1d8d5ec5d091b460562ffd545e4a5e91d17d4a0db7ebe733be34ed374377bd
+size 4733389

tokenizer_config.json

@@ -1,12 +1,19 @@
 {
-  "add_prefix_space": false,
   "backend": "tokenizers",
-  "bos_token": "<|endoftext|>",
+  "bos_token": "<|startoftext|>",
   "clean_up_tokenization_spaces": false,
-  "eos_token": "<|endoftext|>",
+  "eos_token": "<|im_end|>",
   "is_local": false,
+  "legacy": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
   "model_max_length": 1000000000000000019884624838656,
-  "pad_token": "<|padding|>",
+  "pad_token": "<|pad|>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
   "tokenizer_class": "TokenizersBackend",
-  "unk_token": "<|endoftext|>"
+  "use_default_system_prompt": false,
+  "use_fast": true
 }

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:26a0e94f7a02ee2ce084ebf0698e80bc5004f06d8b188a3891c9033114516d03
+oid sha256:4f58299e731f70094b4ca534665893befd6f679e7b0ef988525dd21ec51db615
 size 5329