Model save

README.md

@@ -1,536 +1,84 @@
 ---
-language: en
-license: apache-2.0
+library_name: transformers
 tags:
-- causal-lm
-- research
-- bf16
-- attention
-- normalization
-- neollm
-datasets:
-- HuggingFaceFW/fineweb-edu
+- generated_from_trainer
+model-index:
+- name: NeoLLM
+  results: []
 ---
 
-# NeoLLM
-
-NeoLLM is a **85.16 M parameter** decoder-only language model trained from scratch on
-[FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) with
-**BF16** compute, completing training in approximately
-**~1h 00m** on **NVIDIA RTX PRO 6000 Blackwell Server Edition**.
-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.
-
-> **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) | 4 |
-| Head dim | 64 |
-| Attention layers | full: 4, ckda: 8 |
-| Intermediate size | 1536 |
-| Vocabulary | LiquidAI/LFM2.5-1.2B-Thinking tokenizer (64,402 tokens) |
-| Context length | 512 tokens |
-| Leviathan | enabled=True; d_seed=128, modes=8, knots=16 |
-| JTok / JTok-M | enabled=False / False; Torch continuous surfaces, experts=5, top_k=2 |
-| DAGFormer | enabled=False; {'predictor_dim': 256, 'predictor_layers': 2, 'predictor_heads': 4, 'local_dim': 128, 'input_source': 'tokens'} |
-| Linear Depth Memory | enabled=False; groups=8, key dimension=8; shared branch history when combined with SiameseNorm |
-| Virtual Head Attention | enabled=False; query/KV heads follow the configured geometry |
-| Copy-Paste LM | enabled=True; causal vocabulary/pointer mixture |
-
-### Parameter breakdown
-
-| Parameter bucket | Count |
-|---|---|
-| **Total parameters** | 85.16M (85,162,728) |
-| **Embedding parameters** (tied) | 32.97M (32,973,824) |
-| **Non-embedding parameters** | 52.19M (52,188,904) |
-| **Effective trainable parameters** | 85.16M (85,162,728) |
-
-> 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 = 52.19M`.
-
-### Integrated techniques
-
-NeoLLM combines architecture modules, optional auxiliary objectives, and
-training-time optimizer/stability components from the following papers.
-
-**Embedding and token representation**
-
-- **Learnable Multipliers** ([arXiv:2601.04890](https://arxiv.org/abs/2601.04890)) — Adds
-  per-row and per-column learnable scalar parameters to selected matrix layers and, when
-  enabled, embeddings.
-- **Leviathan** ([arXiv:2601.22040](https://arxiv.org/abs/2601.22040)) — Optional
-  continuous token embedding generator that can replace the discrete input lookup table.
-- **Leviathan-JTok / JTok-M** ([arXiv:2602.00800](https://arxiv.org/abs/2602.00800)) —
-  Optional continuous adaptation of JTok over the shared Leviathan coordinate. JTok
-  applies one normalized surface as a multiplicative MLP gate; JTok-M routes a
-  Sigmoid+TopK mixture of continuous surfaces and adds the paper's depth-scaled
-  residual. Runtime ``jtok_kernel_backend`` selects the PyTorch reference or the
-  CCE Triton adapter; the JTok surfaces remain outside MXFP8 conversion.
-- **KHRONOS** ([arXiv:2505.13315](https://arxiv.org/abs/2505.13315)) — Kernel/basis
-  reference used by the Leviathan continuous token generator implementation.
-- **Spelling Bee Embeddings** ([arXiv:2601.18030](https://arxiv.org/abs/2601.18030)) —
-  Augments token embeddings with character-level spelling information.
-- **Token Embedding Manifold analysis** ([arXiv:2504.01002](https://arxiv.org/abs/2504.01002)) —
-  Reference motivation for treating token embeddings as structured objects rather than
-  unconstrained lookup rows.
-
-**Attention, positions, and output projection**
-
-- **Virtual Head Attention** ([arXiv:2610.17959](https://arxiv.org/abs/2610.17959)) — Optional
-  group-specific query premix and low-rank head postmix. NeoLLM shares its gated query
-  projection across corresponding head slots in KV groups and applies postmix after the gate; inference weights are
-  not automatically folded into standard GQA projections.
-- **Copy-Paste LM** ([arXiv:2610.19247](https://arxiv.org/abs/2610.19247)) — Optional
-  causal pointer and sink mixed with the vocabulary distribution. CUDA training
-  uses CCE's sparse Copy-Paste kernels; inference returns dense mixed logits.
-- **FAN** ([arXiv:2502.21309](https://arxiv.org/abs/2502.21309)) — Fourier Analysis Networks.
-  A portion of the projection channels are dedicated to periodic cosine/sine features.
-- **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).
-- **Interleaved Head Attention / IHA** ([arXiv:2602.21371](https://arxiv.org/abs/2602.21371)) —
-  Builds pseudo-heads from learned cross-head mixtures to create multiple attention patterns
-  per original head.
-- **REPO** ([arXiv:2512.14391](https://arxiv.org/abs/2512.14391)) — Context re-positioning
-  module that learns contextual position coordinates above a configurable start layer.
-- **GRAPE** ([arXiv:2512.07805](https://arxiv.org/abs/2512.07805)) — Group representational
-  position encoding used by the REPO-GRAPE positional path.
-- **GOAT priors** ([arXiv:2601.15380](https://arxiv.org/abs/2601.15380)) — Optional
-  factorized attention log-prior channels inspired by trainable attention priors.
-- **Hadamard output projection** ([arXiv:2603.08343](https://arxiv.org/abs/2603.08343)) —
-  Replaces dense attention output projection with a structured Hadamard transform plus
-  lightweight scaling.
-
-**Normalization, residual flow, and MLP**
-
-- **SeeDNorm** ([arXiv:2510.22777](https://arxiv.org/abs/2510.22777)) — Applied to Q and K
-  projections. Dynamically rescales normalization from the input's own statistics.
-- **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.
-- **GPAS** ([arXiv:2506.22049](https://arxiv.org/abs/2506.22049)) — Gradient-Preserving
-  Activation Scaling for residual junctions.
-- **PolyNorm** ([arXiv:2602.04902](https://arxiv.org/abs/2602.04902)) — Replaces the standard
-  MLP activation with normalized linear, quadratic, and cubic branches.
-- **SimpleGPT** ([arXiv:2602.01212](https://arxiv.org/abs/2602.01212)) — Second-order
-  geometry-inspired normalization strategy applied inside MLP projections.
-- **StackMemory / STACKTRANS** ([NeurIPS 2025](https://openreview.net/forum?id=2bbDg587uh)) —
-  Optional differentiable hidden-state stack between decoder layers.
-- **Attention Residuals / AttnRes** ([arXiv:2603.15031](https://arxiv.org/abs/2603.15031)) —
-  Optional learned depth-wise aggregation over previous layer outputs or block summaries.
-- **LAUREL** ([arXiv:2411.07501](https://arxiv.org/abs/2411.07501)) — Optional learned
-  augmented residual layer with residual-weight and low-rank variants.
-
-**Training objectives and training-time regularizers**
-
-- **Cut Cross Entropy** ([Apple repository](https://github.com/apple/ml-cross-entropy)) —
-  Memory-efficient next-token loss that avoids materializing the full token-by-vocabulary
-  logits tensor. NeoLLM remains compatible with the upstream package when the extensions
-  below are disabled.
-- **MiLe Loss** ([arXiv:2310.19531](https://arxiv.org/abs/2310.19531)) — Optional detached,
-  mean-normalized predictive-entropy weighting of token losses, implemented inside the
-  extended CCE path.
-- **Output Embedding Centering / mu-loss**
-  ([arXiv:2601.02031](https://arxiv.org/abs/2601.02031)) — Optional
-  `lambda * ||mean(output_embeddings)||^2` regularizer for output-logit stability.
-- **MEAP** ([arXiv:2502.07490](https://arxiv.org/abs/2502.07490)) — Optional training-only
-  input corruption that masks a fixed fraction of eligible tokens while preserving clean
-  next-token labels, causal attention, and the inference path.
-- **TWEO** ([arXiv:2511.23225](https://arxiv.org/abs/2511.23225)) — Optional
-  Transformers Without Extreme Outliers activation regularizer for FP8/low-bit-friendly
-  training.
-- **NITP** ([arXiv:2605.24956](https://arxiv.org/abs/2605.24956)) — Optional Next Implicit
-  Token Prediction auxiliary objective using shallow-layer implicit token targets and a
-  cosine loss.
-- **NextLat** ([arXiv:2511.05963](https://arxiv.org/abs/2511.05963)) — Optional next-latent
-  prediction objective using latent dynamics, Smooth L1 supervision, and frozen-head KL.
-
-### Optional extended-CCE configuration
-
-| Feature | Enabled | Value |
-|---|---:|---:|
-| MiLe Loss | True | gamma=1.0 |
-| mu-loss | True | lambda=0.0001 |
-| MEAP | True | ratio=0.15 |
-
-Copy-Paste uses CCE sparse mixture statistics for CUDA training and computes mu-loss in Torch. MEAP and patch training also use the extended cut-cross-entropy package.
-[`Kitsunp/ml-cross-entropy`](https://github.com/Kitsunp/ml-cross-entropy) supplies the extended kernels.
-Copy-Paste exposes three compact
-scalars: unweighted NTP cross entropy, the MiLe reweighting delta, and the mu-loss penalty.
-Their sum reconstructs `ntp_loss` exactly. MEAP reports its eligible and selected counts from
-the masking kernel; the trainer logs the selected count and exact fraction. Copy-Paste diagnostics use sparse source-token candidates and CCE vocabulary reductions.
-
-**Optimizer and training stability**
-
-- **Adam-family optimizer mode** — Active variant: **AdEMAMix**.
-  **AdEMAMix** ([arXiv:2409.03137](https://arxiv.org/abs/2409.03137)) contributes the optional
-  beta3 slow-memory branch when `ADEMAMIX=True`. **DeltaMomentum**
-  ([arXiv:2608.19491](https://arxiv.org/abs/2608.19491)) optionally replaces only the fast beta1
-  first-moment update on hidden Linear weights when `DELTA=True`; beta2 remains Adam's
-  ordinary squared-gradient EMA in every mode.
-- **BF16 stochastic rounding** ([arXiv:2502.20566](https://arxiv.org/abs/2502.20566)) —
-  Optional master-free BF16 model update with stochastic rather than nearest rounding; the same
-  write-back path is shared by AdamW, DeltaAdamW, AdEMAMix, and DeltaAdEMAMix.
-- **Cautious Weight Decay** ([arXiv:2510.12402](https://arxiv.org/abs/2510.12402)) —
-  Sign-selective weight decay variant used by the custom optimizer logic.
-- **Correction of Decoupled Weight Decay** ([arXiv:2512.08217](https://arxiv.org/abs/2512.08217)) —
-  Adapts decoupled weight decay during learning-rate decay.
-- **AdamHD** ([arXiv:2511.14721](https://arxiv.org/abs/2511.14721)) —
-  Decoupled Huber decay regularization reference used by the optimizer.
-- **Gradient conditioning** — CAST active; history-adaptive spectral/norm policies by optimizer group; 5208 committed updates.
-  References: [CAST](https://arxiv.org/abs/2610.19761),
-  [GradientStabilizer](https://arxiv.org/abs/2502.17055).
----
-
-## Training
-
-| Setting | Value |
-|---|---|
-| Dataset | FineWeb-Edu (sample-10BT) |
-| Tokens seen | ~0.51B (15,625 steps × batch 64 × length 512) |
-| Precision | BF16 model state (bfloat16: 85,157,752 elements, float32: 4,976 elements) + AdEMAMix FP32 master active; stochastic rounding disabled |
-| Optimizer | AdEMAMix |
-| Gradient transform | CAST active; history-adaptive spectral/norm policies by optimizer group; 5208 committed updates |
-| Learning rate | 1e-03 with 4,000-step linear warmup |
-| Patch transition | Reset optimizer/conditioner at step 10,417; per-phase LR warmup; restarted AdEMAMix warmup fitted to each phase |
-| Weight decay | coefficients [0.005, 0.01, 0.1]; WD 5/6 groups; CWD 2; Huber 0; Correction 0 |
-| Training time | ~1h 00m |
-| Hardware | NVIDIA RTX PRO 6000 Blackwell Server Edition |
-
+<!-- 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. -->
 
-### Training curve — total objective
-
-Total includes MiLe when enabled, plus all active auxiliary objectives and regularizers. Without MiLe means model total minus its MiLe delta on the same forward; this is not a separate training run without MiLe.
-
-| Step | Train total | Val total | Train without MiLe | Val without MiLe |
-|---|---|---|---|---|
-| 1 | 12.1569 | — | 12.1569 | — |
-| 250 | 7.43977 | — | 7.43193 | — |
-| 500 | 7.26973 | — | 7.2544 | — |
-| 750 | 7.19437 | — | 7.17487 | — |
-| 1,000 | 7.14684 | — | 7.12439 | — |
-| 1,250 | 7.11603 | — | 7.09145 | — |
-| 1,500 | 7.09486 | — | 7.06766 | — |
-| 1,750 | 7.06701 | — | 7.03783 | — |
-| 2,000 | 7.04502 | — | 7.01332 | — |
-| 2,250 | 7.03388 | — | 7.0007 | — |
-| 2,500 | 7.00954 | — | 6.97522 | — |
-| 2,750 | 6.99406 | — | 6.95859 | — |
-| 3,000 | 6.97407 | — | 6.93746 | — |
-| 3,250 | 6.95565 | — | 6.91757 | — |
-| 3,500 | 6.93633 | — | 6.89698 | — |
-| 3,750 | 6.91393 | — | 6.8733 | — |
-| 4,000 | 6.89179 | — | 6.85102 | — |
-| 4,250 | 6.87827 | — | 6.83644 | — |
-| 4,500 | 6.85901 | — | 6.81621 | — |
-| 4,750 | 6.83507 | — | 6.79141 | — |
-| 5,000 | 6.81285 | 6.72223 | 6.76761 | 6.65072 |
-| 5,250 | 6.79736 | — | 6.75163 | — |
-| 5,500 | 6.78718 | — | 6.7406 | — |
-| 5,750 | 6.76926 | — | 6.72092 | — |
-| 6,000 | 6.75714 | — | 6.70777 | — |
-| 6,250 | 6.74706 | — | 6.69821 | — |
-| 6,500 | 6.73372 | — | 6.68512 | — |
-| 6,750 | 6.72918 | — | 6.67894 | — |
-| 7,000 | 6.71321 | — | 6.66257 | — |
-| 7,250 | 6.709 | — | 6.65742 | — |
-| 7,500 | 6.70526 | — | 6.65312 | — |
-| 7,750 | 6.69506 | — | 6.64233 | — |
-| 8,000 | 6.68961 | — | 6.63594 | — |
-| 8,250 | 6.6791 | — | 6.62425 | — |
-| 8,500 | 6.67233 | — | 6.61803 | — |
-| 8,750 | 6.66846 | — | 6.61245 | — |
-| 9,000 | 6.65872 | — | 6.60286 | — |
-| 9,250 | 6.64321 | — | 6.58543 | — |
-| 9,500 | 6.63519 | — | 6.57696 | — |
-| 9,750 | 6.63123 | — | 6.57312 | — |
-| 10,000 | 6.61406 | 6.49322 | 6.55421 | 6.39542 |
-| 10,250 | 6.60986 | — | 6.54967 | — |
-| 10,500 | 6.35387 | — | 6.23053 | — |
-| 10,750 | 5.52453 | — | 5.18023 | — |
-| 11,000 | 5.41604 | — | 5.04308 | — |
-| 11,250 | 5.35595 | — | 4.96565 | — |
-| 11,500 | 5.3127 | — | 4.9082 | — |
-| 11,750 | 5.2762 | — | 4.86209 | — |
-| 12,000 | 5.249 | — | 4.82771 | — |
-| 12,250 | 5.22035 | — | 4.79211 | — |
-| 12,500 | 5.19179 | — | 4.75853 | — |
-| 12,750 | 5.18146 | — | 4.73926 | — |
-| 13,000 | 5.15687 | — | 4.71174 | — |
-| 13,250 | 5.14689 | — | 4.69688 | — |
-| 13,500 | 5.12962 | — | 4.67441 | — |
-| 13,750 | 5.13156 | — | 4.67378 | — |
-| 14,000 | 5.11436 | — | 4.65388 | — |
-| 14,250 | 5.10853 | — | 4.64635 | — |
-| 14,500 | 5.10456 | — | 4.63878 | — |
-| 14,750 | 5.08745 | — | 4.61515 | — |
-| 15,000 | 5.0795 | 4.72568 | 4.60567 | 4.15717 |
-| 15,250 | 5.04957 | — | 4.56981 | — |
-| 15,500 | 5.03392 | — | 4.55032 | — |
-| 15,625 | 5.01443 | 4.67434 | 4.52845 | 4.09541 |
-
-### Cross-entropy with and without MiLe
-
-CE without MiLe is unweighted CE; CE with MiLe adds the reweighting delta. Both exclude mu and auxiliary terms. `ntp_loss` includes mu when enabled.
-
-| Step | Train CE without MiLe | Val CE without MiLe | Train CE with MiLe | Val CE with MiLe |
-|---|---|---|---|---|
-| 1 | 11.1701 | — | 11.1701 | — |
-| 250 | 7.29727 | — | 7.30511 | — |
-| 500 | 6.9752 | — | 6.99054 | — |
-| 750 | 6.85647 | — | 6.87596 | — |
-| 1,000 | 6.78162 | — | 6.80407 | — |
-| 1,250 | 6.72941 | — | 6.75399 | — |
-| 1,500 | 6.68879 | — | 6.71599 | — |
-| 1,750 | 6.64708 | — | 6.67627 | — |
-| 2,000 | 6.61682 | — | 6.64852 | — |
-| 2,250 | 6.59735 | — | 6.63052 | — |
-| 2,500 | 6.56865 | — | 6.60297 | — |
-| 2,750 | 6.5461 | — | 6.58158 | — |
-| 3,000 | 6.52786 | — | 6.56447 | — |
-| 3,250 | 6.51125 | — | 6.54932 | — |
-| 3,500 | 6.49461 | — | 6.53397 | — |
-| 3,750 | 6.47505 | — | 6.51568 | — |
-| 4,000 | 6.45678 | — | 6.49755 | — |
-| 4,250 | 6.44566 | — | 6.48749 | — |
-| 4,500 | 6.43062 | — | 6.47341 | — |
-| 4,750 | 6.41224 | — | 6.4559 | — |
-| 5,000 | 6.39302 | 6.10403 | 6.43826 | 6.17555 |
-| 5,250 | 6.38207 | — | 6.4278 | — |
-| 5,500 | 6.37265 | — | 6.41923 | — |
-| 5,750 | 6.3573 | — | 6.40565 | — |
-| 6,000 | 6.34639 | — | 6.39577 | — |
-| 6,250 | 6.33966 | — | 6.38852 | — |
-| 6,500 | 6.32792 | — | 6.37652 | — |
-| 6,750 | 6.32322 | — | 6.37345 | — |
-| 7,000 | 6.30764 | — | 6.35829 | — |
-| 7,250 | 6.30256 | — | 6.35414 | — |
-| 7,500 | 6.29889 | — | 6.35103 | — |
-| 7,750 | 6.28792 | — | 6.34064 | — |
-| 8,000 | 6.28031 | — | 6.33398 | — |
-| 8,250 | 6.26715 | — | 6.322 | — |
-| 8,500 | 6.26087 | — | 6.31517 | — |
-| 8,750 | 6.25325 | — | 6.30926 | — |
-| 9,000 | 6.2418 | — | 6.29765 | — |
-| 9,250 | 6.22361 | — | 6.28139 | — |
-| 9,500 | 6.21424 | — | 6.27247 | — |
-| 9,750 | 6.20976 | — | 6.26787 | — |
-| 10,000 | 6.19095 | 5.83395 | 6.2508 | 5.93174 |
-| 10,250 | 6.1862 | — | 6.24638 | — |
-| 10,500 | 5.84885 | — | 5.9722 | — |
-| 10,750 | 4.83343 | — | 5.17773 | — |
-| 11,000 | 4.70232 | — | 5.07528 | — |
-| 11,250 | 4.62258 | — | 5.01289 | — |
-| 11,500 | 4.56063 | — | 4.96513 | — |
-| 11,750 | 4.50957 | — | 4.92368 | — |
-| 12,000 | 4.47011 | — | 4.8914 | — |
-| 12,250 | 4.43025 | — | 4.85849 | — |
-| 12,500 | 4.39278 | — | 4.82604 | — |
-| 12,750 | 4.36963 | — | 4.81182 | — |
-| 13,000 | 4.33854 | — | 4.78368 | — |
-| 13,250 | 4.32028 | — | 4.77029 | — |
-| 13,500 | 4.29482 | — | 4.75004 | — |
-| 13,750 | 4.29068 | — | 4.74846 | — |
-| 14,000 | 4.26889 | — | 4.72937 | — |
-| 14,250 | 4.25928 | — | 4.72146 | — |
-| 14,500 | 4.24998 | — | 4.71576 | — |
-| 14,750 | 4.22652 | — | 4.69881 | — |
-| 15,000 | 4.21672 | 3.78368 | 4.69055 | 4.35219 |
-| 15,250 | 4.18105 | — | 4.6608 | — |
-| 15,500 | 4.16243 | — | 4.64603 | — |
-| 15,625 | 4.14067 | 3.72466 | 4.62666 | 4.30359 |
-
-### Observed loss components — weighted contributions
-
-Only contributions actually added to the objective are shown; raw values and subcomponents remain in training metrics. `regularization_loss` is the subtotal of mu, weighted TWeo, JTok-M and StackMemory entropy, excluding NITP/NextLat; do not add it again to those components. JTok-M auxiliary loss is training-only.
-
-| Step | Phase | mu | TWeo | NITP | Reg. subtotal |
-|---|---|---|---|---|---|
-| 1 | train | 3.29888e-10 | 8.48626e-07 | 0.986843 | 8.48956e-07 |
-| 250 | train | 4.14796e-06 | 0.000991383 | 0.133666 | 0.000995531 |
-| 500 | train | 9.26922e-06 | 0.000522309 | 0.27866 | 0.000531578 |
-| 750 | train | 1.54994e-05 | 0.000563869 | 0.317831 | 0.000579369 |
-| 1,000 | train | 2.29899e-05 | 0.000667879 | 0.34208 | 0.000690869 |
-| 1,250 | train | 3.11781e-05 | 0.000880633 | 0.361128 | 0.000911811 |
-| 1,500 | train | 3.90434e-05 | 0.0010814 | 0.377748 | 0.00112044 |
-| 1,750 | train | 4.62851e-05 | 0.00120793 | 0.389492 | 0.00125421 |
-| 2,000 | train | 5.29924e-05 | 0.0013658 | 0.395081 | 0.00141879 |
-| 2,250 | train | 5.94789e-05 | 0.00148332 | 0.401808 | 0.0015428 |
-| 2,500 | train | 6.67849e-05 | 0.00154828 | 0.404959 | 0.00161507 |
-| 2,750 | train | 7.49795e-05 | 0.00155874 | 0.410846 | 0.00163372 |
-| 3,000 | train | 8.46922e-05 | 0.00156715 | 0.407948 | 0.00165184 |
-| 3,250 | train | 9.67422e-05 | 0.00153115 | 0.404699 | 0.0016279 |
-| 3,500 | train | 0.000110971 | 0.00144551 | 0.400806 | 0.00155648 |
-| 3,750 | train | 0.000130353 | 0.00136114 | 0.396762 | 0.00149149 |
-| 4,000 | train | 0.000154362 | 0.00130591 | 0.392777 | 0.00146027 |
-| 4,250 | train | 0.000184102 | 0.00123947 | 0.389356 | 0.00142357 |
-| 4,500 | train | 0.000221222 | 0.00121635 | 0.38416 | 0.00143757 |
-| 4,750 | train | 0.000270394 | 0.00112581 | 0.377766 | 0.0013962 |
-| 5,000 | train | 0.000329293 | 0.00104372 | 0.373213 | 0.00137302 |
-| 5,000 | validation | 0.000360502 | 0.00300473 | 0.543321 | 0.00336524 |
-| 5,250 | train | 0.000392398 | 0.000965072 | 0.368199 | 0.00135747 |
-| 5,500 | train | 0.00045961 | 0.000925218 | 0.366562 | 0.00138483 |
-| 5,750 | train | 0.000528098 | 0.000901479 | 0.362185 | 0.00142958 |
-| 6,000 | train | 0.000591771 | 0.000816563 | 0.359964 | 0.00140833 |
-| 6,250 | train | 0.000653964 | 0.000787942 | 0.357107 | 0.00144191 |
-| 6,500 | train | 0.000715662 | 0.000744753 | 0.355735 | 0.00146041 |
-| 6,750 | train | 0.000772059 | 0.000724611 | 0.354228 | 0.00149667 |
-| 7,000 | train | 0.000824695 | 0.0007098 | 0.353387 | 0.0015345 |
-| 7,250 | train | 0.000869477 | 0.000695535 | 0.353298 | 0.00156501 |
-| 7,500 | train | 0.000910874 | 0.000696223 | 0.35262 | 0.0016071 |
-| 7,750 | train | 0.000952081 | 0.000657018 | 0.3528 | 0.0016091 |
-| 8,000 | train | 0.00099915 | 0.000646211 | 0.353984 | 0.00164536 |
-| 8,250 | train | 0.00104221 | 0.000617805 | 0.355436 | 0.00166002 |
-| 8,500 | train | 0.00107749 | 0.000568957 | 0.355507 | 0.00164644 |
-| 8,750 | train | 0.00110638 | 0.000522036 | 0.357573 | 0.00162841 |
-| 9,000 | train | 0.00113453 | 0.000480492 | 0.359452 | 0.00161502 |
-| 9,250 | train | 0.00115607 | 0.000435784 | 0.360226 | 0.00159185 |
-| 9,500 | train | 0.00117275 | 0.000402017 | 0.361148 | 0.00157476 |
-| 9,750 | train | 0.00118449 | 0.000366701 | 0.361806 | 0.00155119 |
-| 10,000 | train | 0.00118936 | 0.000332726 | 0.361733 | 0.00152209 |
-| 10,000 | validation | 0.00119075 | 0.00078579 | 0.559497 | 0.00197654 |
-| 10,250 | train | 0.00119141 | 0.00030158 | 0.361981 | 0.00149299 |
-| 10,500 | train | 0.00119632 | 0.000360923 | 0.380119 | 0.00155724 |
-| 10,750 | train | 0.0012154 | 0.000378246 | 0.345206 | 0.00159365 |
-| 11,000 | train | 0.00120116 | 0.000295311 | 0.339263 | 0.00149647 |
-| 11,250 | train | 0.00117451 | 0.000256025 | 0.341638 | 0.00143054 |
-| 11,500 | train | 0.00114554 | 0.000232495 | 0.346194 | 0.00137804 |
-| 11,750 | train | 0.00112226 | 0.00021073 | 0.351185 | 0.00133299 |
-| 12,000 | train | 0.00109982 | 0.000196519 | 0.356304 | 0.00129634 |
-| 12,250 | train | 0.00107239 | 0.000192012 | 0.360594 | 0.0012644 |
-| 12,500 | train | 0.001045 | 0.000192727 | 0.364515 | 0.00123773 |
-| 12,750 | train | 0.00102371 | 0.000192792 | 0.368415 | 0.0012165 |
-| 13,000 | train | 0.00100721 | 0.000202993 | 0.371987 | 0.0012102 |
-| 13,250 | train | 0.000990148 | 0.000218566 | 0.375393 | 0.00120871 |
-| 13,500 | train | 0.0009695 | 0.000237501 | 0.378378 | 0.001207 |
-| 13,750 | train | 0.00094896 | 0.000262072 | 0.381884 | 0.00121103 |
-| 14,000 | train | 0.000940227 | 0.000291677 | 0.383756 | 0.0012319 |
-| 14,250 | train | 0.00093996 | 0.000329635 | 0.385801 | 0.00126959 |
-| 14,500 | train | 0.000946027 | 0.00038465 | 0.387467 | 0.00133068 |
-| 14,750 | train | 0.000957708 | 0.000446461 | 0.387233 | 0.00140417 |
-| 15,000 | train | 0.000971831 | 0.00044881 | 0.387528 | 0.00142064 |
-| 15,000 | validation | 0.000978731 | 0.000483302 | 0.372026 | 0.00146203 |
-| 15,250 | train | 0.000985148 | 0.000454519 | 0.387328 | 0.00143967 |
-| 15,500 | train | 0.000994999 | 0.000437405 | 0.386456 | 0.0014324 |
-| 15,625 | train | 0.000999827 | 0.000417227 | 0.386356 | 0.00141705 |
-| 15,625 | validation | 0.00100091 | 0.000434212 | 0.369312 | 0.00143513 |
-
-Tables use exact logged steps; missing values are — without carrying values forward. Repeated evaluations at one step use the latest logged value. Components average batch scalars weighted by examples; CE/PPL token metrics retain their valid-target weighting. Weight decay/CWD act in the optimizer, outside these forward losses.
-
----
-
-## Limitations
-
-- **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.
-
----
-
-## References
-
-All papers whose techniques are integrated into NeoLLM's architecture,
-training objective, or training stack:
-
-| Area | Technique | Paper title | Reference |
-|---|---|---|---|
-| Embeddings | Learnable Multipliers | Freeing the Scale of Language Model Matrix Layers | [arXiv:2601.04890](https://arxiv.org/abs/2601.04890) |
-| Embeddings | Leviathan | A Separable Architecture for Continuous Token Representation in Language Models | [arXiv:2601.22040](https://arxiv.org/abs/2601.22040) |
-| Embeddings / residual modulation | JTok / JTok-M | JTok: On Token Embedding as another Axis of Scaling Law via Joint Token Self-modulation | [arXiv:2602.00800](https://arxiv.org/abs/2602.00800) |
-| Embeddings | KHRONOS | KHRONOS: a Kernel-Based Neural Architecture for Rapid, Resource-Efficient Scientific Computation | [arXiv:2505.13315](https://arxiv.org/abs/2505.13315) |
-| Embeddings | Spelling Bee | Spelling Bee Embeddings for Language Modeling | [arXiv:2601.18030](https://arxiv.org/abs/2601.18030) |
-| Embeddings | Token embedding analysis | Token Embeddings Violate the Manifold Hypothesis | [arXiv:2504.01002](https://arxiv.org/abs/2504.01002) |
-| Attention / positions | FAN | Fourier Analysis Networks | [arXiv:2502.21309](https://arxiv.org/abs/2502.21309) |
-| Attention / positions | MEA | Explicit Multi-head Attention for Inter-head Interaction in Large Language Models | [arXiv:2601.19611](https://arxiv.org/abs/2601.19611) |
-| Attention / positions | LUCID | Attention with Preconditioned Representations | [arXiv:2602.10410](https://arxiv.org/abs/2602.10410) |
-| Attention / positions | Affine-Scaled Attention | Affine-Scaled Attention: Towards Flexible and Stable Transformer Attention | [arXiv:2602.23057](https://arxiv.org/abs/2602.23057) |
-| Attention / positions | XSA | Exclusive Self Attention | [arXiv:2603.09078](https://arxiv.org/abs/2603.09078) |
-| Attention / positions | Directional Routing | Directional Routing in Transformers | [arXiv:2603.14923](https://arxiv.org/abs/2603.14923) |
-| Attention / positions | Gated Attention | Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free | [arXiv:2505.06708](https://arxiv.org/abs/2505.06708) |
-| Attention / positions | Momentum Attention | Momentum Attention | [arXiv:2411.03884](https://arxiv.org/abs/2411.03884) |
-| Attention / positions | IHA | Interleaved Head Attention | [arXiv:2602.21371](https://arxiv.org/abs/2602.21371) |
-| Attention / positions | REPO | Language Models with Context Re-Positioning | [arXiv:2512.14391](https://arxiv.org/abs/2512.14391) |
-| Attention / positions | GRAPE | Group Representational Position Encoding | [arXiv:2512.07805](https://arxiv.org/abs/2512.07805) |
-| Attention / positions | GOAT priors | You Need Better Attention Priors | [arXiv:2601.15380](https://arxiv.org/abs/2601.15380) |
-| Attention / positions | Hadamard o_proj | Rethinking Attention Output Projection: Structured Hadamard Transforms for Efficient Transformers | [arXiv:2603.08343](https://arxiv.org/abs/2603.08343) |
-| Residual / normalization | SeeDNorm | Self-Rescaled Dynamic Normalization | [arXiv:2510.22777](https://arxiv.org/abs/2510.22777) |
-| Residual / normalization | LNS | The Curse of Depth in LLMs | [arXiv:2502.05795](https://arxiv.org/abs/2502.05795) |
-| Residual / normalization | GPAS | Gradient-Preserving Activation Scaling | [arXiv:2506.22049](https://arxiv.org/abs/2506.22049) |
-| Residual / normalization | PolyNorm | PolyNorm / PolyCom | [arXiv:2602.04902](https://arxiv.org/abs/2602.04902) |
-| Residual / normalization | SimpleGPT | SimpleGPT | [arXiv:2602.01212](https://arxiv.org/abs/2602.01212) |
-| Residual / normalization | StackMemory / STACKTRANS | Recursive Transformer: Boosting Reasoning Ability with State Stack | [NeurIPS 2025](https://openreview.net/forum?id=2bbDg587uh) |
-| Residual / normalization | Attention Residuals | Attention Residuals | [arXiv:2603.15031](https://arxiv.org/abs/2603.15031) |
-| Residual / normalization | Linear Depth Memory | Depth Memory | [arXiv:2610.18809](https://arxiv.org/abs/2610.18809) |
-| Residual / normalization | LAUREL | LAUREL: Learned Augmented Residual Layer | [arXiv:2411.07501](https://arxiv.org/abs/2411.07501) |
-| Objectives | TWEO | Transformers Without Extreme Outliers Enables FP8 Training And Quantization For Dummies | [arXiv:2511.23225](https://arxiv.org/abs/2511.23225) |
-| Objectives | NITP | Next Implicit Token Prediction for LLM Pre-training | [arXiv:2605.24956](https://arxiv.org/abs/2605.24956) |
-| Objectives | NextLat | Next-Latent Prediction Transformers Learn Compact World Models | [arXiv:2511.05963](https://arxiv.org/abs/2511.05963) |
-| Optimizer / training | AdEMAMix | The AdEMAMix Optimizer: Better, Faster, Older | [arXiv:2409.03137](https://arxiv.org/abs/2409.03137) |
-| Optimizer / training | DeltaMomentum | DeltaMomentum: A Key-Value based Anisotropic Momentum Update via Delta Rule | [arXiv:2608.19491](https://arxiv.org/abs/2608.19491) |
-| Optimizer / training | BF16 + SR | Stochastic Rounding for LLM Training: Theory and Practice | [arXiv:2502.20566](https://arxiv.org/abs/2502.20566) |
-| Optimizer / training | CWD | Cautious Weight Decay | [arXiv:2510.12402](https://arxiv.org/abs/2510.12402) |
-| Optimizer / training | WD correction | Correction of Decoupled Weight Decay | [arXiv:2512.08217](https://arxiv.org/abs/2512.08217) |
-| Optimizer / training | AdamHD | AdamHD: Decoupled Huber Decay Regularization for Language Model Pre-Training | [arXiv:2511.14721](https://arxiv.org/abs/2511.14721) |
-| Optimizer / training | GradientStabilizer | GradientStabilizer | [arXiv:2502.17055](https://arxiv.org/abs/2502.17055) |
-
----
-
-## Citation
-
-```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}
-}
-```
-
----
-
-## Author
-
-[@Kyokopom](https://x.com/Kyokopom) on X
-
----
-
-## License
+# NeoLLM
 
-Apache 2.0
+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: 4.6689
+- Ce: 3.6325
+- Valid Targets: 2083162.0
+- Ppl: 37.8089
+- Mile Loss: 4.2281
+- Mile Ppl: 68.5884
+- Top1 Accuracy: 0.3089
+- Top3 Accuracy: 0.4925
+- Topk Valid Targets: 2083162.0
+- Copy Gate Mean: 0.2791
+- Copy Sink Mean: 0.6099
+- Copy Effective Gate Mean: 0.1439
+- Copy Target Responsibility Mean: 0.1414
+- Copy Target In Context Fraction: 0.4607
+- Copy Valid Targets: 2083162.0
+- Ntp Loss: 4.2294
+- Ntp Ce Unweighted: 3.6331
+- Mile Reweighting Delta: 0.5955
+- Mu Loss: 0.0008
+- Tweo Loss: 4.8147
+- Nitp Loss: 0.3914
+- Total Model Loss: 4.6689
+- Tweo Weighted Loss: 0.0481
+- Nitp Weighted Loss: 0.3914
+- Auxiliary Loss: 0.4395
+- Regularization Loss: 0.0489
+
+## Model description
+
+More information needed
+
+## Intended uses & limitations
+
+More information needed
+
+## Training and evaluation data
+
+More information needed
+
+## Training procedure
+
+### Training hyperparameters
+
+The following hyperparameters were used during training:
+- learning_rate: 0.001
+- 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
+- num_epochs: 1
+
+### Training results
+
+| Training Loss | Epoch | Step  | Validation Loss | Ce     | Valid Targets | Ppl     | Mile Loss | Mile Ppl | Top1 Accuracy | Top3 Accuracy | Topk Valid Targets | Copy Gate Mean | Copy Sink Mean | Copy Effective Gate Mean | Copy Target Responsibility Mean | Copy Target In Context Fraction | Copy Valid Targets | Ntp Loss | Ntp Ce Unweighted | Mile Reweighting Delta | Mu Loss | Tweo Loss | Nitp Loss | Total Model Loss | Tweo Weighted Loss | Nitp Weighted Loss | Auxiliary Loss | Regularization Loss |
+|:-------------:|:-----:|:-----:|:---------------:|:------:|:-------------:|:-------:|:---------:|:--------:|:-------------:|:-------------:|:------------------:|:--------------:|:--------------:|:------------------------:|:-------------------------------:|:-------------------------------:|:------------------:|:--------:|:-----------------:|:----------------------:|:-------:|:---------:|:---------:|:----------------:|:------------------:|:------------------:|:--------------:|:-------------------:|
+| 5.5014        | 0.32  | 5000  | 5.1422          | 4.0948 | 2083162.0     | 60.0244 | 4.6077    | 100.2542 | 0.2594        | 0.4360        | 2083162.0          | 0.3215         | 0.5617         | 0.1674                   | 0.1650                          | 0.4607                          | 2083162.0          | 4.6084   | 4.0952            | 0.5129                 | 0.0002  | 5.6657    | 0.4772    | 5.1422           | 0.0567             | 0.4772             | 0.5338         | 0.0568              |
+| 5.2629        | 0.64  | 10000 | 4.8937          | 3.8387 | 2083162.0     | 46.4664 | 4.3937    | 80.9401  | 0.2866        | 0.4673        | 2083162.0          | 0.2638         | 0.5579         | 0.1464                   | 0.1483                          | 0.4607                          | 2083162.0          | 4.3947   | 3.8392            | 0.5549                 | 0.0006  | 5.4333    | 0.4446    | 4.8937           | 0.0543             | 0.4446             | 0.4989         | 0.0549              |
+| 5.0697        | 0.96  | 15000 | 4.6966          | 3.6587 | 2083162.0     | 38.8125 | 4.2554    | 70.4849  | 0.3060        | 0.4893        | 2083162.0          | 0.2770         | 0.6103         | 0.1435                   | 0.1423                          | 0.4607                          | 2083162.0          | 4.2566   | 3.6593            | 0.5966                 | 0.0008  | 4.8820    | 0.3911    | 4.6966           | 0.0488             | 0.3911             | 0.4399         | 0.0496              |
+| 5.0386        | 1.0   | 15625 | 4.6689          | 3.6325 | 2083162.0     | 37.8089 | 4.2281    | 68.5884  | 0.3089        | 0.4925        | 2083162.0          | 0.2791         | 0.6099         | 0.1439                   | 0.1414                          | 0.4607                          | 2083162.0          | 4.2294   | 3.6331            | 0.5955                 | 0.0008  | 4.8147    | 0.3914    | 4.6689           | 0.0481             | 0.3914             | 0.4395         | 0.0489              |
+
+
+### Framework versions
+
+- Transformers 5.22.1
+- Pytorch 2.14.1+cu130
+- Datasets 4.8.5
+- Tokenizers 0.23.3

config.json

@@ -159,7 +159,7 @@
   "use_ldm": false,
   "use_learnable_multipliers": true,
   "use_liger_kernel": false,
-  "use_linear_attention": true,
+  "use_linear_attention": false,
   "use_lns": true,
   "use_lucid_attention": false,
   "use_mea_attention": false,

configuration_neollm.py

@@ -356,9 +356,9 @@ class NeoLLMConfig(PretrainedConfig):
             normalized-input mechanism used by the strongest SiameseNorm
             variant.
         siamese_depth_scaling (:obj:`bool`, *optional*, defaults to ``True``):
-            Scale the residual update added to the bounded SiameseNorm stream
-            by ``1 / sqrt(2 * (layer_idx + 1))``. The unbounded stream receives
-            the unscaled shared update.
+            Store the X stream as ``sqrt(2 * (layer_idx + 1)) * RMSNorm(X) +
+            update`` after each shared sublayer, following the official
+            SiameseNorm post/pre implementation. Y receives the shared update.
         siamese_attn_x_scale_init (:obj:`float`, *optional*, defaults to ``1.0``):
             Initial value for the learnable elementwise scale applied to the
             bounded stream in the SiameseNorm attention input.
@@ -928,9 +928,9 @@ class NeoLLMConfig(PretrainedConfig):
         # ── CKDA / GDN2 linear attention ──────────────────────────────────
         # Keyword-only to preserve positional compatibility with old configs.
         *,
-        use_linear_attention=True,             # None: infer; False: disable; True: require linear layers
+        use_linear_attention=False,             # None: infer; False: disable; True: require linear layers
         linear_attention_type="ckda",          # "none", "ckda", or "gdn2"
-        linear_attention_full_attention_every=3,  # 3 linear + 1 full with 12 layers
+        linear_attention_full_attention_every=3,  # 2 linear + 1 full with 12 layers
         linear_attention_layer_types=None,      # Optional per-layer list overrides type and interval
         linear_attention_num_heads=None,        # None: inherit num_attention_heads
         linear_attention_head_dim=None,         # None: inherit the model's head_dim

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:fcfdc55c933f180c5313e598e158ca51eaacf4a45307e73365976527a9d36588
-size 170384136
+oid sha256:66dc2f6bbe5587cd8054d940bd728671738fff4227dda29a0989d4d7311681f8
+size 170287752

modeling_neollm.py

@@ -8031,6 +8031,12 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         self.siamesenorm_kernel_backend = getattr(config, "siamesenorm_kernel_backend", "reference")
         if self.siamesenorm_kernel_backend not in {"reference", "triton", "cute"}:
             raise ValueError("siamesenorm_kernel_backend must be reference, triton or cute.")
+        if self.use_siamesenorm and self.siamesenorm_kernel_backend != "reference":
+            if layer_idx == 0:
+                logger.warning("SiameseNorm sublayers use PyTorch: the selected custom backend "
+                               "does not implement the corrected residual topology. "
+                               "The final stream merge retains the selected backend.")
+            self.siamesenorm_kernel_backend = "reference"
         self.siamese_normalized_input = bool(
             getattr(config, "siamese_normalized_input", True)
         )
@@ -8040,19 +8046,18 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         self.siamese_attn_x_scale_init = float(
             getattr(config, "siamese_attn_x_scale_init", 1.0)
         )
-        siamese_stream_scale = (
+        siamese_depth_factor = (
             1.0
             if not self.siamese_depth_scaling
-            else 1.0 / math.sqrt(2.0 * float(layer_idx + 1))
+            else math.sqrt(2.0 * float(layer_idx + 1))
         )
         # A tensor buffer keeps the depth-dependent value out of Python control
         # flow during forward.  Computing it from ``self.layer_idx`` inside a
         # Dynamo resume frame specialized that frame once per decoder layer.
-        # Non-persistent preserves checkpoint compatibility: the value is fully
-        # determined by layer_idx and the config on every construction.
+        # The factor is determined by the layer index and configuration.
         self.register_buffer(
-            "_siamese_stream_scale_value",
-            torch.tensor(siamese_stream_scale, dtype=torch.float32),
+            "_siamese_depth_factor",
+            torch.tensor(siamese_depth_factor, dtype=torch.float32),
             persistent=True,
         )
         # Controls only the first pre-attention normalisation applied directly
@@ -8104,13 +8109,13 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             # constructed only when the Siamese topology is active, so no
             # inactive RMSNorm pre-norm modules remain in the graph.
             self.siamese_attn_x_norm = nn.RMSNorm(
-                config.hidden_size, eps=config.rms_norm_eps
+                config.hidden_size, eps=config.rms_norm_eps, elementwise_affine=False
             )
             self.siamese_attn_y_norm = nn.RMSNorm(
                 config.hidden_size, eps=config.rms_norm_eps
             )
             self.siamese_mlp_x_norm = nn.RMSNorm(
-                config.hidden_size, eps=config.rms_norm_eps
+                config.hidden_size, eps=config.rms_norm_eps, elementwise_affine=False
             )
             self.siamese_mlp_y_norm = nn.RMSNorm(
                 config.hidden_size, eps=config.rms_norm_eps
@@ -8132,6 +8137,9 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
                     dtype=torch.float32,
                 )
             )
+            self.siamese_mlp_x_scale = nn.Parameter(
+                torch.ones(config.hidden_size, dtype=torch.float32)
+            )
         else:
             self.input_layernorm = (
                 _make_norm(
@@ -8152,6 +8160,7 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             self.siamese_attn_input_norm = None
             self.siamese_mlp_input_norm = None
             self.siamese_attn_x_scale = None
+            self.siamese_mlp_x_scale = None
 
         self.lns_attn = LNS(layer_idx) if self.use_lns else None
         self.lns_mlp = LNS(layer_idx) if self.use_lns else None
@@ -8452,8 +8461,8 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             self._last_dag_r_dynamics = packed.detach().clone()
         return h, coefficients, residual, y_residual
 
-    def _siamese_stream_scale(self, ref: torch.Tensor) -> torch.Tensor:
-        return self._siamese_stream_scale_value.to(
+    def _siamese_depth_scale(self, ref: torch.Tensor) -> torch.Tensor:
+        return self._siamese_depth_factor.to(
             device=ref.device,
             dtype=ref.dtype,
         )
@@ -8507,7 +8516,6 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         # parameters.  All Siamese normalization modules are RMSNorm by
         # construction; no dynamic-normalization branch remains.
         # ── Attention shared block ────────────────────────────────────────
-        norm_backend = getattr(self, "siamesenorm_kernel_backend", "reference")
         if self.use_msar:
             branch = 2 * self.layer_idx
             msar_sources, msar_partial = _msar_start_branch(
@@ -8516,29 +8524,16 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
                 branch, self.msar_block_branches, msar_valid, msar_position)
         else:
             msar_attn = None
-        if norm_backend == "reference":
-            x_attn_norm = self.siamese_attn_x_norm(x_states)
-            y_attn_norm = self.siamese_attn_y_norm(y_states)
-            x_scale = self.siamese_attn_x_scale.to(
-                dtype=x_attn_norm.dtype, device=x_attn_norm.device
-            )
-            h_attn = x_scale * x_attn_norm + y_attn_norm
-        else:
-            from cut_cross_entropy.siamesenorm import siamesenorm
-            h_attn = siamesenorm(
-                x_states, y_states,
-                self.siamese_attn_x_norm.weight, self.siamese_attn_y_norm.weight,
-                self.siamese_attn_input_norm.weight
-                if not (self.use_msar or self.use_ldm) and self.siamese_attn_input_norm is not None else None,
-                self.siamese_attn_x_scale.to(dtype=x_states.dtype, device=x_states.device),
-                eps=self.siamese_attn_x_norm.eps, backend=norm_backend,
-            )
+        x_attn_norm = self.siamese_attn_x_norm(x_states)
+        y_attn_norm = self.siamese_attn_y_norm(y_states)
+        x_scale = self.siamese_attn_x_scale.to(dtype=x_attn_norm.dtype, device=x_attn_norm.device)
+        h_attn = x_scale * x_attn_norm + y_attn_norm
         if self.use_msar:
             h_attn = h_attn + msar_attn
         if self.use_ldm:
             ldm_memory, ldm_covariance, ldm_attn = self.ldm_attn_read(h_attn, ldm_state, ldm_valid)
             h_attn = h_attn + ldm_attn
-        if self.siamese_attn_input_norm is not None and (norm_backend == "reference" or self.use_msar or self.use_ldm):
+        if self.siamese_attn_input_norm is not None:
             h_attn = self.siamese_attn_input_norm(h_attn)
 
         stack_metrics = None
@@ -8570,12 +8565,12 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             **kwargs,
         )
 
-        stream_scale = self._siamese_stream_scale(attn_out)
-        if dag_rx is not None:
-            x_states = x_states + dag_rx
-            y_states = y_states + dag_ry
-        x_after_attn = x_states + stream_scale * attn_out
+        depth_factor = self._siamese_depth_scale(attn_out)
+        x_after_attn = depth_factor * x_attn_norm + attn_out
         y_after_attn = y_states + attn_out
+        if dag_rx is not None:
+            x_after_attn = x_after_attn + dag_rx
+            y_after_attn = y_after_attn + dag_ry
         if self.use_ldm:
             ldm_state = (ldm_memory, ldm_covariance,
                          *self.ldm_attn_write((x_after_attn + y_after_attn) * .5,
@@ -8595,37 +8590,21 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
                 branch, self.msar_block_branches, msar_valid, msar_position)
         else:
             msar_mlp = None
-        if norm_backend == "reference":
-            x_mlp_norm = self.siamese_mlp_x_norm(x_after_attn)
-            y_mlp_norm = self.siamese_mlp_y_norm(y_after_attn)
-            h_mlp = x_mlp_norm + y_mlp_norm
-        else:
-            h_mlp = siamesenorm(
-                x_after_attn, y_after_attn,
-                self.siamese_mlp_x_norm.weight, self.siamese_mlp_y_norm.weight,
-                self.siamese_mlp_input_norm.weight
-                if not (self.use_msar or self.use_ldm) and self.siamese_mlp_input_norm is not None else None,
-                eps=self.siamese_mlp_x_norm.eps, backend=norm_backend,
-                post_scale=(
-                    self.lns_mlp.scale.to(device=x_after_attn.device, dtype=x_after_attn.dtype)
-                    if self.use_lns and not (self.use_msar or self.use_ldm) else None
-                ),
-                out_dtype=(
-                    torch.get_autocast_dtype("cuda")
-                    if not (self.use_msar or self.use_ldm) and torch.is_autocast_enabled("cuda") else None
-                ),
-            )
+        x_mlp_norm = self.siamese_mlp_x_norm(x_after_attn)
+        y_mlp_norm = self.siamese_mlp_y_norm(y_after_attn)
+        mlp_scale = self.siamese_mlp_x_scale.to(dtype=x_mlp_norm.dtype, device=x_mlp_norm.device)
+        h_mlp = mlp_scale * x_mlp_norm + y_mlp_norm
         if self.use_msar:
             h_mlp = h_mlp + msar_mlp
         if self.use_ldm:
             ldm_memory, ldm_covariance, ldm_mlp = self.ldm_mlp_read(h_mlp, ldm_state, ldm_valid)
             h_mlp = h_mlp + ldm_mlp
-        if self.siamese_mlp_input_norm is not None and (norm_backend == "reference" or self.use_msar or self.use_ldm):
+        if self.siamese_mlp_input_norm is not None:
             h_mlp = self.siamese_mlp_input_norm(h_mlp)
 
         # Only this consumer starts directly with an autocast Linear. Attention
         # also exposes native states to positional, router and memory consumers.
-        h_lns2 = self.lns_mlp(h_mlp) if self.use_lns and (norm_backend == "reference" or self.use_msar or self.use_ldm) else h_mlp
+        h_lns2 = self.lns_mlp(h_mlp) if self.use_lns else h_mlp
         delta_m = self.mlp(h_lns2)
 
         jtok_aux_stats = None
@@ -8644,7 +8623,7 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         else:
             shared_update = delta_m
 
-        x_next = x_after_attn + stream_scale * shared_update
+        x_next = depth_factor * x_mlp_norm + shared_update
         y_next = y_after_attn + shared_update
         if self.use_ldm:
             ldm_state = (ldm_memory, ldm_covariance,
@@ -9297,18 +9276,17 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
             and requested_dtype.is_floating_point
         ):
             for module in model.modules():
-                siamese_scale = getattr(module, "siamese_attn_x_scale", None)
-                if siamese_scale is not None:
-                    siamese_scale.data = siamese_scale.data.to(
-                        dtype=requested_dtype
-                    )
+                for name in ("siamese_attn_x_scale", "siamese_mlp_x_scale"):
+                    siamese_scale = getattr(module, name, None)
+                    if siamese_scale is not None:
+                        siamese_scale.data = siamese_scale.data.to(dtype=requested_dtype)
                 stream_scale = getattr(
                     module,
-                    "_siamese_stream_scale_value",
+                    "_siamese_depth_factor",
                     None,
                 )
                 if stream_scale is not None:
-                    module._buffers["_siamese_stream_scale_value"] = (
+                    module._buffers["_siamese_depth_factor"] = (
                         stream_scale.to(dtype=requested_dtype)
                     )
         if return_loading_info:
@@ -9475,6 +9453,8 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
                 and module.siamese_attn_x_scale is not None
             ):
                 module.siamese_attn_x_scale.data.fill_(module.siamese_attn_x_scale_init)
+            if getattr(module, "siamese_mlp_x_scale", None) is not None:
+                module.siamese_mlp_x_scale.data.fill_(1.0)
             # AttnRes pseudo-queries: MUST be initialized to zero.
             # Zero initialization ensures uniform attention weights at step 0
             # (softmax of zeros is uniform), making AttnRes equivalent to a
@@ -12520,9 +12500,10 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     if gpas is not None:
                         metrics[prefix + f"gpas_{slot}_forward_gain"] = (1-F.silu(gpas.alpha.detach().float())).reshape(())
                 if layer.use_siamesenorm:
-                    scale = layer.siamese_attn_x_scale.detach().float()
-                    metrics[prefix + "siamese_attn_x_scale_mean"] = scale.mean()
-                    metrics[prefix + "siamese_attn_x_scale_rms"] = scale.square().mean().sqrt()
+                    for slot in ("attn", "mlp"):
+                        scale = getattr(layer, f"siamese_{slot}_x_scale").detach().float()
+                        metrics[prefix + f"siamese_{slot}_x_scale_mean"] = scale.mean()
+                        metrics[prefix + f"siamese_{slot}_x_scale_rms"] = scale.square().mean().sqrt()
                 if layer.layer_idx > 0 and layer.self_attn.use_fan_residual:
                     for name in ("lambda_1", "lambda_2"):
                         metrics[prefix + "fan_" + name] = getattr(layer.self_attn, name).detach().float().reshape(())
@@ -12574,6 +12555,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 "laurel": ("laurel_rw_attn", "laurel_rw_mlp", "laurel_lr_A_attn",
                            "laurel_lr_B_attn", "laurel_lr_A_mlp", "laurel_lr_B_mlp"),
                 "siamese": ("siamese_attn_x_scale",),
+                "siamese_mlp": ("siamese_mlp_x_scale",),
                 "stack": ("stack_memory",),
                 "gpas": ("gpas_attn", "gpas_mlp"),
             }

optimizer_precision.json

@@ -8,9 +8,9 @@
   "delta_enabled": false,
   "delta_eta": 0.37,
   "fp32_master_active": true,
-  "fp32_master_elements": 85162728,
-  "fp32_master_groups": 7,
-  "fp32_master_tensors": 406,
+  "fp32_master_elements": 85119928,
+  "fp32_master_groups": 6,
+  "fp32_master_tensors": 368,
   "fp32_passthrough_elements": 0,
   "mode": "fp32_master",
   "mxfp8_active": false,
@@ -19,12 +19,12 @@
   "mxfp8_recipe": null,
   "native_precision_elements": 0,
   "optimizer_variant": "AdEMAMix",
-  "ordinary_linear_layers": 154,
+  "ordinary_linear_layers": 129,
   "parameter_elements_by_dtype": {
-    "bfloat16": 85157752,
-    "float32": 4976
+    "bfloat16": 85117752,
+    "float32": 2176
   },
-  "parameter_storage_bytes": 170335408,
+  "parameter_storage_bytes": 170244208,
   "runtime_precision_mode": "BF16_FP32_MASTER",
   "stochastic_round_seed": 0
 }

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:177c8a895c89df3c73884f8275f81a6a7412bf270fa73920044b3a3c23164a60
+oid sha256:9874bd7fc0bda2eb6503079473525ee5fe8948e0ee4bb947a748e8326822ac5b
 size 5329