Model save

README.md

@@ -1,532 +1,78 @@
 ---
-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.09 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
-**~0h 47m** 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 |
-| 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'} |
-| 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.09M (85,093,304) |
-| **Embedding parameters** (tied) | 32.97M (32,973,824) |
-| **Non-embedding parameters** | 52.12M (52,119,480) |
-| **Effective trainable parameters** | 85.09M (85,093,304) |
-
-> 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.12M`.
-
-### 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; 15625 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,091,128 elements, float32: 2,176 elements) + AdEMAMix FP32 master active; stochastic rounding disabled |
-| Optimizer | AdEMAMix |
-| Gradient transform | CAST active; history-adaptive spectral/norm policies by optimizer group; 15625 committed updates |
-| Learning rate | 1e-03 with 4,000-step linear warmup |
-| Weight decay | coefficients [0.005, 0.01, 0.1]; WD 5/5 groups; CWD 2; Huber 0; Correction 0 |
-| Training time | ~0h 47m |
-| 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.2845 | — | 12.2845 | — |
-| 250 | 7.51328 | — | 7.50593 | — |
-| 500 | 7.37632 | — | 7.36168 | — |
-| 750 | 7.33534 | — | 7.31743 | — |
-| 1,000 | 7.2956 | — | 7.27509 | — |
-| 1,250 | 7.25757 | — | 7.2354 | — |
-| 1,500 | 7.22776 | — | 7.20332 | — |
-| 1,750 | 7.18971 | — | 7.16376 | — |
-| 2,000 | 7.16679 | — | 7.13879 | — |
-| 2,250 | 7.15075 | — | 7.12142 | — |
-| 2,500 | 7.1182 | — | 7.08786 | — |
-| 2,750 | 7.09542 | — | 7.064 | — |
-| 3,000 | 7.07595 | — | 7.0437 | — |
-| 3,250 | 7.05828 | — | 7.02471 | — |
-| 3,500 | 7.04374 | — | 7.00896 | — |
-| 3,750 | 7.03011 | — | 6.99452 | — |
-| 4,000 | 7.01765 | — | 6.98176 | — |
-| 4,250 | 7.00604 | — | 6.96907 | — |
-| 4,500 | 6.99175 | — | 6.95386 | — |
-| 4,750 | 6.97323 | — | 6.93439 | — |
-| 5,000 | 6.95435 | 7.03772 | 6.91428 | 6.97998 |
-| 5,250 | 6.94098 | — | 6.90037 | — |
-| 5,500 | 6.9301 | — | 6.88864 | — |
-| 5,750 | 6.91726 | — | 6.87426 | — |
-| 6,000 | 6.905 | — | 6.86125 | — |
-| 6,250 | 6.89508 | — | 6.85136 | — |
-| 6,500 | 6.8811 | — | 6.8378 | — |
-| 6,750 | 6.87753 | — | 6.8327 | — |
-| 7,000 | 6.86382 | — | 6.81881 | — |
-| 7,250 | 6.85904 | — | 6.81314 | — |
-| 7,500 | 6.85565 | — | 6.8089 | — |
-| 7,750 | 6.84594 | — | 6.79925 | — |
-| 8,000 | 6.8399 | — | 6.79186 | — |
-| 8,250 | 6.8271 | — | 6.77829 | — |
-| 8,500 | 6.82443 | — | 6.77634 | — |
-| 8,750 | 6.8251 | — | 6.77586 | — |
-| 9,000 | 6.81961 | — | 6.77061 | — |
-| 9,250 | 6.80826 | — | 6.75788 | — |
-| 9,500 | 6.80517 | — | 6.75465 | — |
-| 9,750 | 6.805 | — | 6.7549 | — |
-| 10,000 | 6.79291 | 6.90652 | 6.74131 | 6.83695 |
-| 10,250 | 6.79435 | — | 6.74302 | — |
-| 10,500 | 6.57804 | — | 6.46818 | — |
-| 10,750 | 5.64887 | — | 5.30678 | — |
-| 11,000 | 5.50961 | — | 5.1322 | — |
-| 11,250 | 5.44243 | — | 5.04573 | — |
-| 11,500 | 5.39135 | — | 4.97991 | — |
-| 11,750 | 5.3537 | — | 4.93408 | — |
-| 12,000 | 5.32865 | — | 4.90204 | — |
-| 12,250 | 5.29837 | — | 4.86656 | — |
-| 12,500 | 5.26806 | — | 4.83252 | — |
-| 12,750 | 5.25518 | — | 4.81263 | — |
-| 13,000 | 5.22838 | — | 4.78307 | — |
-| 13,250 | 5.2163 | — | 4.7663 | — |
-| 13,500 | 5.1947 | — | 4.74032 | — |
-| 13,750 | 5.19089 | — | 4.73445 | — |
-| 14,000 | 5.16889 | — | 4.71043 | — |
-| 14,250 | 5.1559 | — | 4.69624 | — |
-| 14,500 | 5.14548 | — | 4.68216 | — |
-| 14,750 | 5.13203 | — | 4.66433 | — |
-| 15,000 | 5.13145 | 4.7568 | 4.66355 | 4.18981 |
-| 15,250 | 5.11204 | — | 4.64183 | — |
-| 15,500 | 5.10473 | — | 4.63283 | — |
-| 15,625 | 5.09438 | 4.7291 | 4.62103 | 4.15709 |
-
-### 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.2377 | — | 11.2377 | — |
-| 250 | 7.33083 | — | 7.33818 | — |
-| 500 | 7.01992 | — | 7.03457 | — |
-| 750 | 6.9117 | — | 6.92962 | — |
-| 1,000 | 6.84122 | — | 6.86173 | — |
-| 1,250 | 6.78967 | — | 6.81184 | — |
-| 1,500 | 6.75051 | — | 6.77496 | — |
-| 1,750 | 6.71397 | — | 6.73992 | — |
-| 2,000 | 6.68534 | — | 6.71334 | — |
-| 2,250 | 6.66963 | — | 6.69896 | — |
-| 2,500 | 6.64192 | — | 6.67227 | — |
-| 2,750 | 6.621 | — | 6.65242 | — |
-| 3,000 | 6.60442 | — | 6.63668 | — |
-| 3,250 | 6.59009 | — | 6.62366 | — |
-| 3,500 | 6.5758 | — | 6.61058 | — |
-| 3,750 | 6.55762 | — | 6.59321 | — |
-| 4,000 | 6.54081 | — | 6.5767 | — |
-| 4,250 | 6.52894 | — | 6.56591 | — |
-| 4,500 | 6.51391 | — | 6.5518 | — |
-| 4,750 | 6.49645 | — | 6.53529 | — |
-| 5,000 | 6.4769 | 6.35905 | 6.51696 | 6.41679 |
-| 5,250 | 6.46543 | — | 6.50604 | — |
-| 5,500 | 6.45592 | — | 6.49738 | — |
-| 5,750 | 6.441 | — | 6.484 | — |
-| 6,000 | 6.43025 | — | 6.474 | — |
-| 6,250 | 6.42147 | — | 6.4652 | — |
-| 6,500 | 6.40953 | — | 6.45284 | — |
-| 6,750 | 6.40453 | — | 6.44936 | — |
-| 7,000 | 6.39054 | — | 6.43555 | — |
-| 7,250 | 6.38482 | — | 6.43072 | — |
-| 7,500 | 6.3804 | — | 6.42715 | — |
-| 7,750 | 6.36999 | — | 6.41668 | — |
-| 8,000 | 6.36302 | — | 6.41106 | — |
-| 8,250 | 6.35259 | — | 6.4014 | — |
-| 8,500 | 6.35024 | — | 6.39833 | — |
-| 8,750 | 6.34786 | — | 6.3971 | — |
-| 9,000 | 6.34222 | — | 6.39122 | — |
-| 9,250 | 6.32834 | — | 6.37873 | — |
-| 9,500 | 6.32368 | — | 6.37419 | — |
-| 9,750 | 6.32497 | — | 6.37507 | — |
-| 10,000 | 6.31065 | 6.20464 | 6.36225 | 6.27421 |
-| 10,250 | 6.31157 | — | 6.3629 | — |
-| 10,500 | 6.01485 | — | 6.12471 | — |
-| 10,750 | 4.89594 | — | 5.23804 | — |
-| 11,000 | 4.72479 | — | 5.1022 | — |
-| 11,250 | 4.63385 | — | 5.03055 | — |
-| 11,500 | 4.56533 | — | 4.97677 | — |
-| 11,750 | 4.5158 | — | 4.93542 | — |
-| 12,000 | 4.48125 | — | 4.90787 | — |
-| 12,250 | 4.44365 | — | 4.87546 | — |
-| 12,500 | 4.40877 | — | 4.84431 | — |
-| 12,750 | 4.38846 | — | 4.83101 | — |
-| 13,000 | 4.35794 | — | 4.80324 | — |
-| 13,250 | 4.33989 | — | 4.78989 | — |
-| 13,500 | 4.31309 | — | 4.76746 | — |
-| 13,750 | 4.30621 | — | 4.76265 | — |
-| 14,000 | 4.28204 | — | 4.7405 | — |
-| 14,250 | 4.26676 | — | 4.72641 | — |
-| 14,500 | 4.2524 | — | 4.71572 | — |
-| 14,750 | 4.23404 | — | 4.70174 | — |
-| 15,000 | 4.23243 | 3.78875 | 4.70032 | 4.35574 |
-| 15,250 | 4.21065 | — | 4.68087 | — |
-| 15,500 | 4.20151 | — | 4.67341 | — |
-| 15,625 | 4.18959 | 3.75747 | 4.66293 | 4.32948 |
-
-### 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.16275e-10 | 8.85299e-07 | 1.0468 | 8.85615e-07 |
-| 250 | train | 4.01168e-06 | 0.00101028 | 0.174095 | 0.00101429 |
-| 500 | train | 8.93938e-06 | 0.000425533 | 0.341321 | 0.000434472 |
-| 750 | train | 1.49341e-05 | 0.000574217 | 0.405132 | 0.000589151 |
-| 1,000 | train | 2.21405e-05 | 0.000717058 | 0.433133 | 0.000739198 |
-| 1,250 | train | 3.00225e-05 | 0.000947325 | 0.444758 | 0.000977348 |
-| 1,500 | train | 3.7833e-05 | 0.00116823 | 0.451597 | 0.00120607 |
-| 1,750 | train | 4.49507e-05 | 0.00147651 | 0.448272 | 0.00152146 |
-| 2,000 | train | 5.20222e-05 | 0.00178034 | 0.451619 | 0.00183236 |
-| 2,250 | train | 5.97428e-05 | 0.00206258 | 0.449671 | 0.00212232 |
-| 2,500 | train | 6.78635e-05 | 0.00220622 | 0.443663 | 0.00227409 |
-| 2,750 | train | 7.61504e-05 | 0.00231168 | 0.440618 | 0.00238783 |
-| 3,000 | train | 8.58517e-05 | 0.00233333 | 0.436857 | 0.00241919 |
-| 3,250 | train | 9.77149e-05 | 0.00236563 | 0.432168 | 0.00246334 |
-| 3,500 | train | 0.000110827 | 0.00231476 | 0.430735 | 0.00242559 |
-| 3,750 | train | 0.000126875 | 0.00226564 | 0.434512 | 0.00239251 |
-| 4,000 | train | 0.000145417 | 0.00211525 | 0.438694 | 0.00226067 |
-| 4,250 | train | 0.000167628 | 0.00196029 | 0.437996 | 0.00212792 |
-| 4,500 | train | 0.000194782 | 0.00182893 | 0.437925 | 0.00202371 |
-| 4,750 | train | 0.000227027 | 0.00170113 | 0.436016 | 0.00192816 |
-| 5,000 | train | 0.000263738 | 0.00154546 | 0.435576 | 0.00180919 |
-| 5,000 | validation | 0.000283719 | 0.0060813 | 0.614562 | 0.00636502 |
-| 5,250 | train | 0.000303911 | 0.00140913 | 0.433229 | 0.00171304 |
-| 5,500 | train | 0.000345929 | 0.00130875 | 0.431072 | 0.00165468 |
-| 5,750 | train | 0.000391956 | 0.00119768 | 0.431669 | 0.00158963 |
-| 6,000 | train | 0.000435046 | 0.00113317 | 0.429428 | 0.00156821 |
-| 6,250 | train | 0.000476276 | 0.00107261 | 0.428334 | 0.00154889 |
-| 6,500 | train | 0.000520947 | 0.000988251 | 0.426753 | 0.0015092 |
-| 6,750 | train | 0.000561467 | 0.000925999 | 0.426688 | 0.00148747 |
-| 7,000 | train | 0.000600642 | 0.000889937 | 0.426782 | 0.00149058 |
-| 7,250 | train | 0.000640212 | 0.000860261 | 0.426819 | 0.00150047 |
-| 7,500 | train | 0.000679096 | 0.000852769 | 0.426972 | 0.00153187 |
-| 7,750 | train | 0.000715778 | 0.000802992 | 0.427743 | 0.00151877 |
-| 8,000 | train | 0.000753992 | 0.00078903 | 0.4273 | 0.00154302 |
-| 8,250 | train | 0.000795478 | 0.000771698 | 0.424126 | 0.00156718 |
-| 8,500 | train | 0.000837118 | 0.000758088 | 0.424502 | 0.00159521 |
-| 8,750 | train | 0.000873988 | 0.00075022 | 0.426376 | 0.00162421 |
-| 9,000 | train | 0.000910428 | 0.000754263 | 0.426734 | 0.00166469 |
-| 9,250 | train | 0.000944036 | 0.000730808 | 0.427862 | 0.00167484 |
-| 9,500 | train | 0.000975691 | 0.00071411 | 0.429285 | 0.0016898 |
-| 9,750 | train | 0.00100823 | 0.000682743 | 0.428235 | 0.00169098 |
-| 10,000 | train | 0.00104292 | 0.000664123 | 0.428943 | 0.00170704 |
-| 10,000 | validation | 0.00106212 | 0.0022162 | 0.629043 | 0.00327832 |
-| 10,250 | train | 0.00108166 | 0.000681578 | 0.429688 | 0.00176324 |
-| 10,500 | train | 0.0011197 | 0.00105885 | 0.451152 | 0.00217854 |
-| 10,750 | train | 0.00114721 | 0.000800759 | 0.408884 | 0.00194796 |
-| 11,000 | train | 0.00114271 | 0.000609513 | 0.405662 | 0.00175222 |
-| 11,250 | train | 0.00112443 | 0.000589441 | 0.410163 | 0.00171387 |
-| 11,500 | train | 0.00110764 | 0.000626262 | 0.412851 | 0.0017339 |
-| 11,750 | train | 0.0011031 | 0.0007034 | 0.416477 | 0.0018065 |
-| 12,000 | train | 0.00110705 | 0.000701879 | 0.418976 | 0.00180893 |
-| 12,250 | train | 0.00111366 | 0.000731281 | 0.42106 | 0.00184494 |
-| 12,500 | train | 0.00112138 | 0.000758367 | 0.421869 | 0.00187974 |
-| 12,750 | train | 0.00113102 | 0.000750959 | 0.422287 | 0.00188198 |
-| 13,000 | train | 0.00114004 | 0.000732814 | 0.423269 | 0.00187286 |
-| 13,250 | train | 0.0011468 | 0.000694571 | 0.424568 | 0.00184138 |
-| 13,500 | train | 0.00115065 | 0.000657363 | 0.425431 | 0.00180801 |
-| 13,750 | train | 0.00115104 | 0.000615922 | 0.426476 | 0.00176697 |
-| 14,000 | train | 0.00115097 | 0.000595539 | 0.42664 | 0.00174651 |
-| 14,250 | train | 0.00115072 | 0.000557399 | 0.427775 | 0.00170812 |
-| 14,500 | train | 0.00114966 | 0.000522873 | 0.428093 | 0.00167254 |
-| 14,750 | train | 0.00115023 | 0.000500229 | 0.428635 | 0.00165046 |
-| 15,000 | train | 0.00115158 | 0.000471158 | 0.429507 | 0.00162274 |
-| 15,000 | validation | 0.00115209 | 0.000515729 | 0.399388 | 0.00166782 |
-| 15,250 | train | 0.00115244 | 0.000452634 | 0.429571 | 0.00160508 |
-| 15,500 | train | 0.00115244 | 0.000431864 | 0.429741 | 0.00158431 |
-| 15,625 | train | 0.00115257 | 0.000415416 | 0.429879 | 0.00156799 |
-| 15,625 | validation | 0.00115263 | 0.000468992 | 0.397996 | 0.00162163 |
-
-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 | 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.9291
+- Ce: 3.9075
+- Valid Targets: 2083162.0
+- Ppl: 49.7745
+- Mile Loss: 4.5021
+- Mile Ppl: 90.2041
+- Top1 Accuracy: 0.3108
+- Top3 Accuracy: 0.4650
+- Topk Valid Targets: 2083162.0
+- Ntp Loss: 4.5034
+- Ntp Ce Unweighted: 3.9080
+- Mile Reweighting Delta: 0.5945
+- Mu Loss: 0.0010
+- Tweo Loss: 0.0668
+- Nitp Loss: 0.4250
+- Total Model Loss: 4.9291
+- Tweo Weighted Loss: 0.0007
+- Nitp Weighted Loss: 0.4250
+- Auxiliary Loss: 0.4257
+- Regularization Loss: 0.0016
+
+## 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 | 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 |
+|:-------------:|:-----:|:-----:|:---------------:|:------:|:-------------:|:--------:|:---------:|:--------:|:-------------:|:-------------:|:------------------:|:--------:|:-----------------:|:----------------------:|:-------:|:---------:|:---------:|:----------------:|:------------------:|:------------------:|:--------------:|:-------------------:|
+| 7.2623        | 0.32  | 5000  | 7.4209          | 6.6778 | 2083162.0     | 794.6030 | 6.7310    | 837.9544 | 0.0784        | 0.1546        | 2083162.0          | 6.7312   | 6.6778            | 0.0531                 | 0.0003  | 0.8364    | 0.6813    | 7.4209           | 0.0084             | 0.6813             | 0.6897         | 0.0087              |
+| 6.9846        | 0.64  | 10000 | 7.1389          | 6.3917 | 2083162.0     | 596.8784 | 6.4607    | 639.4872 | 0.0952        | 0.1738        | 2083162.0          | 6.4614   | 6.3917            | 0.0690                 | 0.0008  | 0.1872    | 0.6756    | 7.1389           | 0.0019             | 0.6756             | 0.6775         | 0.0027              |
+| 5.3718        | 0.96  | 15000 | 4.9829          | 3.9751 | 2083162.0     | 53.2555  | 4.5515    | 94.7726  | 0.3041        | 0.4569        | 2083162.0          | 4.5528   | 3.9756            | 0.5763                 | 0.0009  | 0.0772    | 0.4294    | 4.9829           | 0.0008             | 0.4294             | 0.4301         | 0.0017              |
+| 5.3163        | 1.0   | 15625 | 4.9291          | 3.9075 | 2083162.0     | 49.7745  | 4.5021    | 90.2041  | 0.3108        | 0.4650        | 2083162.0          | 4.5034   | 3.9080            | 0.5945                 | 0.0010  | 0.0668    | 0.4250    | 4.9291           | 0.0007             | 0.4250             | 0.4257         | 0.0016              |
+
+
+### Framework versions
+
+- Transformers 5.22.1
+- Pytorch 2.14.1+cu130
+- Datasets 4.8.5
+- Tokenizers 0.23.3

config.json

@@ -13,6 +13,7 @@
   },
   "bos_token_id": 1,
   "cce_loss_impl": "cce_kahan_full_c",
+  "copy_paste_lexical_keys": false,
   "dagformer": {
     "input_source": "tokens",
     "local_dim": 128,
@@ -56,19 +57,29 @@
   "jtok_num_modes": 4,
   "jtok_top_k": 2,
   "laurel_lr_rank": 32,
+  "ldm_groups": 8,
+  "ldm_hidden_dim": 32,
+  "ldm_key_dim": 8,
   "liger_loss_accum_dtype": "float32",
   "lucid_attention_eps": 1e-06,
   "max_position_embeddings": 512,
   "mea_component_key_value_heads": 4,
   "mea_groupnorm_eps": 1e-06,
   "meap_exclude_last": true,
+  "meap_exclude_mask_from_copy": false,
   "meap_implementation": "triton",
   "meap_mask_ratio": 0.15,
   "meap_mask_token_id": 16,
+  "meap_protect_structural_tokens": false,
   "meap_seed": 42,
   "mile_loss_gamma": 1.0,
   "model_type": "neollm",
   "momentum_gamma": 0.1,
+  "msar_backend": "auto",
+  "msar_block_branches": 4,
+  "msar_kernel_sizes": [
+    4
+  ],
   "mu_loss_lambda": 0.0001,
   "nextlat_dynamics_proj_factor": 1.0,
   "nextlat_horizon": 1,
@@ -113,7 +124,7 @@
   "stack_memory_cache_size": 2048,
   "stack_slots": 16,
   "tie_word_embeddings": false,
-  "transformers_version": "5.21.1",
+  "transformers_version": "5.22.1",
   "tweo_eps": 1e-06,
   "tweo_loss_weight": 0.01,
   "tweo_power": 4.0,
@@ -121,7 +132,7 @@
   "use_affine_scaled_attention": false,
   "use_attn_res": false,
   "use_cache": false,
-  "use_copy_paste": true,
+  "use_copy_paste": false,
   "use_dagformer": false,
   "use_directional_routing": false,
   "use_embedding_input_norm": true,
@@ -136,6 +147,7 @@
   "use_laurel": false,
   "use_laurel_lr": false,
   "use_laurel_rw": false,
+  "use_ldm": false,
   "use_learnable_multipliers": true,
   "use_liger_kernel": false,
   "use_lns": true,
@@ -144,6 +156,7 @@
   "use_meap": true,
   "use_mile_loss": true,
   "use_momentum_attention": true,
+  "use_msar": false,
   "use_mu_loss": true,
   "use_nextlat": false,
   "use_nitp": true,

configuration_neollm.py

@@ -224,12 +224,26 @@ class NeoLLMConfig(PretrainedConfig):
             ``use_xsa=False``.
         use_msar (:obj:`bool`, *optional*, defaults to ``False``):
             Enable Block Multi-Stream Attention Residuals (arXiv:2610.13630).
-            It is exclusive with the other residual topologies, GPAS, FAN
-            residual and DAGFormer. Single-stream depth summaries are retained;
+            It may be combined with SiameseNorm; the hybrid adds an MSAR read
+            of the shared unscaled branch history to the normalized X/Y sum,
+            before its existing optional shared input norm. No extra
+            normalization or user parameter is introduced by the combination.
+            This is an architectural adaptation, not a configuration tested in
+            the MSAR paper. Other residual topologies, GPAS, FAN residual and
+            DAGFormer remain exclusive. Single-stream summaries are retained;
             causal views and independent depth reads are reconstructed per branch.
+        use_ldm (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable Linear Depth Memory (RLS per token and feature group). Group
+            count follows the attention head geometry automatically, and the
+            key and controller dimensions follow the paper's small-width recipe.
+            SiameseNorm can be combined with LDM through a shared history of
+            unscaled branch outputs; this hybrid is an architectural adaptation.
+            Other alternative residual flows remain mutually exclusive.
         msar_block_branches (:obj:`int` or ``None``, optional):
             Number of attention/MLP branches per block. ``None`` derives a
-            size from model depth aiming for roughly eight blocks.
+            size from the total number of attention/MLP branches, aiming for
+            roughly eight contiguous blocks; block boundaries may fall between
+            attention and MLP. An explicit value is preserved as supplied.
         msar_kernel_sizes (:obj:`tuple`, optional):
             Causal depthwise widths; the paper uses ``(4, 8, 12)``.
         msar_backend (:obj:`str`, optional):
@@ -771,8 +785,9 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Attention Residuals (Kimi Team, 2026) ─────────────────────────
         use_attn_res=False,
         attn_res_num_blocks=4,
-        # MSAR: independent residual topology, not combined with SiameseNorm.
+        # MSAR may also be paired with the existing SiameseNorm residual flow.
         use_msar=False,
+        use_ldm=False,
         msar_block_branches=None,  # auto: approximately eight depth blocks
         msar_kernel_sizes=(4, 8, 12),
         msar_backend="auto",  # Triton on CUDA, Torch on CPU; "torch" is a reference
@@ -1104,14 +1119,15 @@ class NeoLLMConfig(PretrainedConfig):
 
         # ── Residual-flow mechanism mutex ─────────────────────────────────
         # AttnRes and LAuReL compete directly with residual topology.
-        # StackMemory may be paired with SiameseNorm because it acts on the
-        # shared block input, while SiameseNorm owns the X/Y stream updates.
+        # StackMemory and MSAR may separately be paired with SiameseNorm. The
+        # hybrid keeps the X/Y updates and reads MSAR's shared branch history.
         _residual_flow_flags = {
             "use_laurel": bool(use_laurel),
             "use_attn_res": bool(use_attn_res),
             "use_stack_memory": bool(use_stack_memory),
             "use_siamesenorm": bool(use_siamesenorm),
             "use_msar": bool(use_msar),
+            "use_ldm": bool(use_ldm),
         }
         _active_residual_flow = [
             name for name, enabled in _residual_flow_flags.items() if enabled
@@ -1122,13 +1138,25 @@ class NeoLLMConfig(PretrainedConfig):
             and not bool(use_laurel)
             and not bool(use_attn_res)
             and not bool(use_msar)
+            and not bool(use_ldm)
         )
-        if len(_active_residual_flow) > 1 and not allow_stack_siamese:
+        allow_msar_siamese = (
+            bool(use_msar) and bool(use_siamesenorm)
+            and not bool(use_stack_memory) and not bool(use_laurel)
+            and not bool(use_attn_res)
+            and not bool(use_ldm)
+        )
+        allow_ldm_siamese = (
+            bool(use_ldm) and bool(use_siamesenorm)
+            and not any((use_stack_memory, use_laurel, use_attn_res, use_msar))
+        )
+        if len(_active_residual_flow) > 1 and not (allow_stack_siamese or allow_msar_siamese or allow_ldm_siamese):
             raise ValueError(
                 "Enable at most one residual-flow mechanism among "
                 "`use_laurel`, `use_attn_res`, `use_stack_memory`, "
-                "`use_siamesenorm`, and `use_msar`, except the explicit "
-                "`use_stack_memory` + `use_siamesenorm` pair. "
+                "`use_siamesenorm`, `use_msar`, and `use_ldm`, except the explicit "
+                "`use_stack_memory` + `use_siamesenorm` and "
+                "`use_msar` + `use_siamesenorm` and `use_ldm` + `use_siamesenorm` pairs. "
                 f"Active flags: {_active_residual_flow}."
             )
 
@@ -1138,9 +1166,15 @@ class NeoLLMConfig(PretrainedConfig):
                 "`use_msar=True` cannot be combined with `use_gpas`, "
                 "`use_fan_residual`, or `use_dagformer`."
             )
+        if use_ldm and (use_gpas or use_fan_residual or use_dagformer):
+            raise ValueError(
+                "`use_ldm=True` cannot be combined with `use_gpas`, "
+                "`use_fan_residual`, or `use_dagformer`."
+            )
         if msar_block_branches is None:
-            # Approximately eight blocks, without tuning for batch or sequence.
-            msar_block_branches = 2 * max(1, math.ceil(num_hidden_layers / 8))
+            # Branch granularity counts attention and MLP independently. Do not
+            # round to whole decoder layers: e.g. 24 branches / 8 = 3.
+            msar_block_branches = max(1, math.ceil(2 * num_hidden_layers / 8))
         if isinstance(msar_block_branches, bool) or not isinstance(msar_block_branches, int) or msar_block_branches < 1:
             raise ValueError("`msar_block_branches` must be a positive integer.")
         if (not isinstance(msar_kernel_sizes, (tuple, list)) or not msar_kernel_sizes
@@ -1382,6 +1416,16 @@ class NeoLLMConfig(PretrainedConfig):
         #   Memory cost: O(num_blocks × batch × seq × hidden) instead of
         #   O(num_layers × batch × seq × hidden).
         self.use_msar = bool(use_msar)
+        self.use_ldm = bool(use_ldm)
+        # Follow the paper's value width = attention head width, even when
+        # num_attention_heads * head_dim differs from hidden_size. For an
+        # incompatible width choose the closest divisor without padding values.
+        divisors = (g for g in range(1, hidden_size + 1) if hidden_size % g == 0)
+        self.ldm_groups = min(divisors,
+                              key=lambda g: (abs(hidden_size // g - head_dim),
+                                             abs(g - num_attention_heads)))
+        self.ldm_key_dim = min(8, hidden_size // self.ldm_groups)
+        self.ldm_hidden_dim = min(32, hidden_size)
         self.msar_block_branches = msar_block_branches
         self.msar_kernel_sizes = tuple(msar_kernel_sizes)
         self.msar_backend = msar_backend

generation_config.json

@@ -7,5 +7,5 @@
   "output_attentions": false,
   "output_hidden_states": false,
   "pad_token_id": 0,
-  "transformers_version": "5.21.1"
+  "transformers_version": "5.22.1"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0dc8e2c81e7e67462b6b6c74089c5f3a25f19051ccaf3f0cedafdc5eb8db79d3
-size 170235776
+oid sha256:30daee49b68ae04dcd13eaf707ee60e1ae2daaf2acb5ff9cf4f0f93975e6fff7
+size 169184776

modeling_neollm.py

@@ -76,6 +76,8 @@ References:
         arXiv:2601.15380.
     StackMemory / STACKTRANS: Zhang et al. (2025). "Recursive Transformer:
         Boosting Reasoning Ability with State Stack." NeurIPS 2025.
+    Linear Depth Memory: González-Martínez et al. (2026). "Depth Memory."
+        arXiv:2610.18809; flexible Triton backend adapted under the MIT license.
     HOLA: Cui (2026). "A Hippocampus for Linear Attention: An Exact Memory
         for What the Recurrent State Forgets." arXiv:2607.02303.
     TWEO: Liang et al. (2026). "Transformers Without Extreme Outliers Enables
@@ -6718,11 +6720,12 @@ def _msar_reference_read(sources, banks, query, valid, position, widths, eps):
 
 if triton is not None:
     @triton.jit
-    def _msar_read_kernel(SOURCES, BANKS, QUERY, VALID, POSITION, OUT, LSE,
+    def _msar_read_kernel(SOURCES, BANKS, QUERY, VALID, POSITION, OUT, LSE, GRAD,
                     T: tl.constexpr, D: tl.constexpr, NS: tl.constexpr,
                     M: tl.constexpr,
-                    WIDTHS: tl.constexpr, EPS: tl.constexpr, BD: tl.constexpr):
-        """Online depth softmax; expanded source views stay in registers."""
+                    WIDTHS: tl.constexpr, EPS: tl.constexpr, BD: tl.constexpr,
+                    DOT_ONLY: tl.constexpr):
+        """Read depth streams, or recompute only the backward scalar <g,z>."""
         row = tl.program_id(0)
         stream = tl.program_id(1)
         batch = row // T
@@ -6734,7 +6737,11 @@ if triton is not None:
         q = tl.load(QUERY + stream * D + c, mask_c, other=0).to(tl.float32)
         running_max = tl.full((), -float('inf'), tl.float32)
         denom = tl.full((), 0., tl.float32)
-        accum = tl.full((BD,), 0., tl.float32)
+        if DOT_ONLY:
+            g = tl.load(GRAD + row * M * D + stream * D + c, mask_c, other=0).to(tl.float32)
+            dot = tl.full((), 0., tl.float32)
+        else:
+            accum = tl.full((BD,), 0., tl.float32)
         for i in tl.static_range(NS):
             value = tl.full((BD,), 0., tl.float32)
             if stream == 0:
@@ -6743,7 +6750,7 @@ if triton is not None:
             offset = 0
             for r in tl.static_range(len(WIDTHS)):
                 if stream == r + 1:
-                    for delta in tl.static_range(WIDTHS[r]):
+                    for delta in tl.static_range(min(WIDTHS[r], T)):
                         past = row - delta
                         in_seq = token >= delta
                         past_valid = tl.load(VALID + past, in_seq, other=0).to(tl.int1)
@@ -6760,91 +6767,141 @@ if triton is not None:
             new_max = tl.maximum(running_max, score)
             old_weight = tl.exp(running_max - new_max)
             new_weight = tl.exp(score - new_max)
-            accum = accum * old_weight + value * new_weight
+            if DOT_ONLY:
+                dot = dot * old_weight + tl.sum(g * value, 0) * new_weight
+            else:
+                accum = accum * old_weight + value * new_weight
             denom = denom * old_weight + new_weight
             running_max = new_max
-        tl.store(OUT + row * M * D + stream * D + c, accum / denom, mask_c)
-        tl.store(LSE + row * M + stream, running_max + tl.log(denom))
+        if DOT_ONLY:
+            tl.store(OUT + row * M + stream, dot / denom)
+        else:
+            tl.store(OUT + row * M * D + stream * D + c, accum / denom, mask_c)
+            tl.store(LSE + row * M + stream, running_max + tl.log(denom))
 
 
     @triton.jit
-    def _msar_vjp_value(SOURCE, BANK, QUERY, VALID, POSITION, LSE, Z, G, DV, QPART,
+    def _msar_vjp_value(SOURCE, BANK, QUERY, VALID, POSITION, LSE, GDOT, G, DV, QPART,
                    ROWS: tl.constexpr, T: tl.constexpr, D: tl.constexpr,
                    M: tl.constexpr,
                    WIDTHS: tl.constexpr, EPS: tl.constexpr, R: tl.constexpr,
-                   BD: tl.constexpr, GROUPS: tl.constexpr):
-        """Per-source VJP through softmax and RMS; reduce query partials by rows."""
+                   BD: tl.constexpr, GROUPS: tl.constexpr, SINGLE_SOURCE: tl.constexpr):
+        """Materialize convolution-view VJPs and their query partials."""
         group = tl.program_id(0)
-        stream = tl.program_id(1)
+        # Identity is handled by _msar_vjp_source: avoid its full DV buffer.
+        stream = tl.program_id(1) + 1
         row = group * R + tl.arange(0, R)
         c = tl.arange(0, BD)
         row_ok = row < ROWS
-        tok = row % T
-        valid = tl.load(VALID + row, row_ok, other=0).to(tl.int1)
-        pos = tl.load(POSITION + row, row_ok, other=0)
-        q = tl.load(QUERY + stream*D+c, c < D, other=0).to(tl.float32)
-        v = tl.full((R, BD), 0., tl.float32)
-        if stream == 0:
-            v = tl.load(SOURCE + row[:,None]*D+c[None,:],
-                        row_ok[:,None] & valid[:,None] & (c[None,:] < D), other=0).to(tl.float32)
-        offset = 0
-        for r in tl.static_range(len(WIDTHS)):
-            if stream == r+1:
-                for delta in tl.static_range(WIDTHS[r]):
-                    past = row - delta
-                    in_seq = row_ok & (tok >= delta)
-                    pv = tl.load(VALID + past, in_seq, other=0).to(tl.int1)
-                    pp = tl.load(POSITION + past, in_seq, other=0)
-                    allowed = in_seq & valid & pv & ((pos-pp) == delta)
-                    x = tl.load(SOURCE + past[:,None]*D+c[None,:],
-                                allowed[:,None] & (c[None,:]<D), other=0).to(tl.float32)
-                    w = tl.load(BANK+(offset+delta)*D+c,
-                                c<D, other=0).to(tl.float32)
-                    v += x*w[None,:]
-            offset += WIDTHS[r]
-        inv = tl.rsqrt(tl.sum(v*v, 1)/D + EPS)
-        qv = tl.sum(v*q[None,:], 1)
         lse = tl.load(LSE+row*M+stream, row_ok, other=0)
-        p = tl.exp(qv*inv-lse)
-        z = tl.load(Z+row[:,None]*M*D+stream*D+c[None,:],
-                    row_ok[:,None] & (c[None,:]<D), other=0).to(tl.float32)
         g = tl.load(G+row[:,None]*M*D+stream*D+c[None,:],
                     row_ok[:,None] & (c[None,:]<D), other=0).to(tl.float32)
-        ds = p*tl.sum(g*(v-z), 1)
-        dv = p[:,None]*g + ds[:,None]*(q[None,:]*inv[:,None] -
-            v*(qv*inv*inv*inv/D)[:,None])
-        tl.store(DV+stream*ROWS*D+row[:,None]*D+c[None,:], dv,
+        if SINGLE_SOURCE:
+            # A finite one-entry softmax is constant: dv=g and dq=0. The
+            # saved LSE and upstream gradient detect nonfinite inputs without
+            # reconstructing any convolution or forming score derivatives.
+            bad_g = tl.sum(((tl.abs(g) >= float('inf')) | (g != g)).to(tl.int32), 1) > 0
+            finite = (tl.abs(lse) < float('inf')) & ~bad_g
+            dv = tl.where(finite[:,None], g, float('nan'))
+            bad_rows = tl.sum((~finite & row_ok).to(tl.int32), 0) > 0
+            # The query VJP is identically zero for a finite one-source
+            # softmax. All D components have the same zero/NaN status, so
+            # store one scalar per row group rather than D identical values.
+            qp = tl.where(bad_rows, float('nan'), 0.)
+        else:
+            tok = row % T
+            valid = tl.load(VALID + row, row_ok, other=0).to(tl.int1)
+            pos = tl.load(POSITION + row, row_ok, other=0)
+            q = tl.load(QUERY + stream*D+c, c < D, other=0).to(tl.float32)
+            v = tl.full((R, BD), 0., tl.float32)
+            offset = 0
+            for r in tl.static_range(len(WIDTHS)):
+                if stream == r+1:
+                    for delta in tl.static_range(min(WIDTHS[r], T)):
+                        past = row - delta
+                        in_seq = row_ok & (tok >= delta)
+                        pv = tl.load(VALID + past, in_seq, other=0).to(tl.int1)
+                        pp = tl.load(POSITION + past, in_seq, other=0)
+                        allowed = in_seq & valid & pv & ((pos-pp) == delta)
+                        x = tl.load(SOURCE + past[:,None]*D+c[None,:],
+                                    allowed[:,None] & (c[None,:]<D), other=0).to(tl.float32)
+                        w = tl.load(BANK+(offset+delta)*D+c,
+                                    c<D, other=0).to(tl.float32)
+                        v += x*w[None,:]
+                offset += WIDTHS[r]
+            inv = tl.rsqrt(tl.sum(v*v, 1)/D + EPS)
+            qv = tl.sum(v*q[None,:], 1)
+            p = tl.exp(qv*inv-lse)
+            gv = tl.sum(g*v, 1)
+            gdot = tl.load(GDOT+row*M+stream, row_ok, other=0)
+            ds = p*(gv-gdot)
+            dv = p[:,None]*g + ds[:,None]*(q[None,:]*inv[:,None] -
+                v*(qv*inv*inv*inv/D)[:,None])
+            qp = tl.sum(ds[:,None]*v*inv[:,None], 0)
+        tl.store(DV+(stream-1)*ROWS*D+row[:,None]*D+c[None,:], dv,
                  row_ok[:,None] & (c[None,:]<D))
-        qp = tl.sum(ds[:,None]*v*inv[:,None], 0)
-        tl.store(QPART+stream*GROUPS*D+group*D+c, qp, c<D)
+        if SINGLE_SOURCE:
+            tl.store(QPART+stream*GROUPS+group, qp)
+        else:
+            tl.store(QPART+stream*GROUPS*D+group*D+c, qp, c<D)
 
 
     @triton.jit
-    def _msar_vjp_source(DV, BANK, VALID, POSITION, DX, ROWS: tl.constexpr,
-                    T: tl.constexpr, D: tl.constexpr,
-                    WIDTHS: tl.constexpr, BLOCK: tl.constexpr):
-        """Transpose the causal depthwise convolutions without atomics."""
-        j = tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
-        row = j//D
-        c = j%D
-        ok = row<ROWS
-        token = row%T
-        valid = tl.load(VALID+row, ok, other=0).to(tl.int1)
-        pos = tl.load(POSITION+row, ok, other=0)
-        dx = tl.load(DV+j, ok & valid, other=0)
+    def _msar_vjp_source(SOURCE, BANK, QUERY, VALID, POSITION, LSE, GDOT, G, DV, QPART, DX,
+                         ROWS: tl.constexpr, T: tl.constexpr, D: tl.constexpr, M: tl.constexpr,
+                         WIDTHS: tl.constexpr, EPS: tl.constexpr, R: tl.constexpr,
+                         BD: tl.constexpr, GROUPS: tl.constexpr, SINGLE_SOURCE: tl.constexpr):
+        """Fuse the identity-view VJP into the convolution input transpose.
+
+        The identity query partial is still reduced in FP32. For a finite
+        one-source softmax its query VJP is exactly zero; nonfinite rows keep
+        the same NaN propagation as the materialized implementation.
+        """
+        group = tl.program_id(0)
+        row = group * R + tl.arange(0, R)
+        c = tl.arange(0, BD)
+        ok = row < ROWS
+        token = row % T
+        valid = tl.load(VALID + row, ok, other=0).to(tl.int1)
+        pos = tl.load(POSITION + row, ok, other=0)
+        lse = tl.load(LSE + row * M, ok, other=0)
+        g = tl.load(G + row[:,None] * M * D + c[None,:],
+                    ok[:,None] & (c[None,:] < D), other=0).to(tl.float32)
+        if SINGLE_SOURCE:
+            bad_g = tl.sum(((tl.abs(g) >= float('inf')) | (g != g)).to(tl.int32), 1) > 0
+            finite = (tl.abs(lse) < float('inf')) & ~bad_g
+            identity_vjp = tl.where(finite[:,None], g, float('nan'))
+            bad_rows = tl.sum((~finite & ok).to(tl.int32), 0) > 0
+            tl.store(QPART + group, tl.where(bad_rows, float('nan'), 0.))
+        else:
+            q = tl.load(QUERY + c, c < D, other=0).to(tl.float32)
+            v = tl.load(SOURCE + row[:,None] * D + c[None,:],
+                        ok[:,None] & valid[:,None] & (c[None,:] < D), other=0).to(tl.float32)
+            inv = tl.rsqrt(tl.sum(v*v, 1) / D + EPS)
+            qv = tl.sum(v*q[None,:], 1)
+            p = tl.exp(qv*inv-lse)
+            gv = tl.sum(g*v, 1)
+            gdot = tl.load(GDOT + row*M, ok, other=0)
+            ds = p*(gv-gdot)
+            identity_vjp = p[:,None]*g + ds[:,None]*(q[None,:]*inv[:,None] -
+                v*(qv*inv*inv*inv/D)[:,None])
+            qp = tl.sum(ds[:,None]*v*inv[:,None], 0)
+            tl.store(QPART + group*D + c, qp, c < D)
+        dx = tl.where(valid[:,None], identity_vjp, 0.)
         offset = 0
         for r in tl.static_range(len(WIDTHS)):
-            for delta in tl.static_range(WIDTHS[r]):
+            for delta in tl.static_range(min(WIDTHS[r], T)):
                 future = row+delta
                 in_seq = ok & ((token+delta)<T)
                 fv = tl.load(VALID+future, in_seq, other=0).to(tl.int1)
                 fp = tl.load(POSITION+future, in_seq, other=0)
                 allowed = in_seq & valid & fv & ((fp-pos)==delta)
-                v = tl.load(DV+(r+1)*ROWS*D+future*D+c, allowed, other=0)
-                w = tl.load(BANK+(offset+delta)*D+c, ok, other=0).to(tl.float32)
-                dx += v*w
+                dv = tl.load(DV+r*ROWS*D+future[:,None]*D+c[None,:],
+                             allowed[:,None] & (c[None,:]<D), other=0)
+                w = tl.load(BANK+(offset+delta)*D+c, c<D, other=0).to(tl.float32)
+                dx += dv*w[None,:]
             offset += WIDTHS[r]
-        tl.store(DX+j, dx, ok)
+        tl.store(DX+row[:,None]*D+c[None,:], dx, ok[:,None] & (c[None,:]<D))
 
 
     @triton.jit
@@ -6861,21 +6918,78 @@ if triton is not None:
         pos = tl.load(POSITION+row, ok, other=0)
         offset = 0
         for r in tl.static_range(len(WIDTHS)):
-            grad = tl.load(DV+(r+1)*ROWS*D+row[:,None]*D+c[None,:],
+            grad = tl.load(DV+r*ROWS*D+row[:,None]*D+c[None,:],
                            ok[:,None] & (c[None,:]<D), other=0).to(tl.float32)
             for delta in tl.static_range(WIDTHS[r]):
-                past = row-delta
-                in_seq = ok & (tok>=delta)
-                pv = tl.load(VALID+past, in_seq, other=0).to(tl.int1)
-                pp = tl.load(POSITION+past, in_seq, other=0)
-                allowed = in_seq & valid & pv & ((pos-pp)==delta)
-                x = tl.load(SOURCE+past[:,None]*D+c[None,:],
-                            allowed[:,None] & (c[None,:]<D), other=0).to(tl.float32)
-                part = tl.sum(grad*x, 0)
+                if delta < T:
+                    past = row-delta
+                    in_seq = ok & (tok>=delta)
+                    pv = tl.load(VALID+past, in_seq, other=0).to(tl.int1)
+                    pp = tl.load(POSITION+past, in_seq, other=0)
+                    allowed = in_seq & valid & pv & ((pos-pp)==delta)
+                    x = tl.load(SOURCE+past[:,None]*D+c[None,:],
+                                allowed[:,None] & (c[None,:]<D), other=0).to(tl.float32)
+                    part = tl.sum(grad*x, 0)
+                else:
+                    # These convolution taps cannot be reached by any token.
+                    # The partial buffer is uninitialized, so write exact zero.
+                    part = tl.full((C,), 0., tl.float32)
                 tl.store(WPART+(offset+delta)*GROUPS*D+group*D+c, part, c<D)
             offset += WIDTHS[r]
 
 
+    @triton.jit
+    def _msar_reduce_query_bank(QPART, WPART, DQ, DW, D: tl.constexpr,
+                                M: tl.constexpr, KSUM: tl.constexpr,
+                                QGROUPS: tl.constexpr, WGROUPS: tl.constexpr,
+                                C: tl.constexpr, REDUCE_ROWS: tl.constexpr,
+                                FIRST_TILE: tl.constexpr, SINGLE_SOURCE: tl.constexpr):
+        """Accumulate query and bank partials in FP32 in one launch.
+
+        A source tile contributes to ``DQ`` across all sources and tiles;
+        ``DW`` belongs to one source and is initialized by its first tile.
+        The single-source query partial is one zero/NaN scalar per row group.
+        Chunked reductions accumulate in FP32, but their order need not be
+        bit-identical to the framework's separate reductions.
+        """
+        pid = tl.program_id(0)
+        column_blocks = tl.cdiv(D, C)
+        if pid < M * column_blocks:
+            q_stream = pid // column_blocks
+            q_col = (pid % column_blocks)*C + tl.arange(0, C)
+            q_offsets = tl.arange(0, REDUCE_ROWS)
+            if SINGLE_SOURCE:
+                q_sum = tl.full((), 0., tl.float32)
+                for q_start in range(tl.cdiv(QGROUPS, REDUCE_ROWS)):
+                    q_row = q_start*REDUCE_ROWS + q_offsets
+                    q_value = tl.load(QPART + q_stream*QGROUPS + q_row,
+                                      q_row < QGROUPS, other=0)
+                    q_sum += tl.sum(q_value, 0)
+            else:
+                q_sum = tl.full((C,), 0., tl.float32)
+                for q_start in range(tl.cdiv(QGROUPS, REDUCE_ROWS)):
+                    q_row = q_start*REDUCE_ROWS + q_offsets
+                    q_value = tl.load(QPART + q_stream*QGROUPS*D + q_row[:,None]*D + q_col[None,:],
+                                      (q_row[:,None] < QGROUPS) & (q_col[None,:] < D), other=0)
+                    q_sum += tl.sum(q_value, 0)
+            previous_q = tl.load(DQ + q_stream*D + q_col, q_col < D, other=0)
+            tl.store(DQ + q_stream*D + q_col, previous_q + q_sum, q_col < D)
+        else:
+            bank_pid = pid - M * column_blocks
+            bank_tap = bank_pid // column_blocks
+            bank_col = (bank_pid % column_blocks)*C + tl.arange(0, C)
+            bank_offsets = tl.arange(0, REDUCE_ROWS)
+            bank_sum = tl.full((C,), 0., tl.float32)
+            for bank_start in range(tl.cdiv(WGROUPS, REDUCE_ROWS)):
+                bank_row = bank_start*REDUCE_ROWS + bank_offsets
+                bank_value = tl.load(WPART + bank_tap*WGROUPS*D + bank_row[:,None]*D + bank_col[None,:],
+                                     (bank_row[:,None] < WGROUPS) & (bank_col[None,:] < D), other=0)
+                bank_sum += tl.sum(bank_value, 0)
+            if not FIRST_TILE:
+                bank_sum += tl.load(DW + bank_tap*D + bank_col, bank_col < D, other=0)
+            tl.store(DW + bank_tap*D + bank_col, bank_sum, bank_col < D)
+
+
 @torch.library.custom_op("neollm_msar::read", mutates_args=())
 def _msar_read_op(sources: list[torch.Tensor], banks: list[torch.Tensor],
                query: torch.Tensor, valid: torch.Tensor, position: torch.Tensor,
@@ -6888,9 +7002,10 @@ def _msar_read_op(sources: list[torch.Tensor], banks: list[torch.Tensor],
     position = position.contiguous()
     result = sources[0].new_empty((b, t, len(widths)+1, d))
     lse = torch.empty((b,t,len(widths)+1),device=result.device,dtype=torch.float32)
-    _msar_read_kernel[(b*t, len(widths)+1)](tuple(sources), tuple(banks), query, valid, position,
-        result, lse, t, d, len(sources), len(widths)+1,
-        tuple(widths), eps, triton.next_power_of_2(d), num_warps=4)
+    if b*t:
+        _msar_read_kernel[(b*t, len(widths)+1)](tuple(sources), tuple(banks), query, valid, position,
+            result, lse, None, t, d, len(sources), len(widths)+1,
+            tuple(widths), eps, triton.next_power_of_2(d), False, num_warps=4)
     return result, lse
 
 
@@ -6906,38 +7021,66 @@ def _msar_backward_op(sources: list[torch.Tensor], banks: list[torch.Tensor],
                   query: torch.Tensor, valid: torch.Tensor, position: torch.Tensor,
                   widths: list[int], eps: float,
                   grad: torch.Tensor, lse: torch.Tensor) -> tuple[list[torch.Tensor], list[torch.Tensor], torch.Tensor]:
-    """Fused VJP with exact FP32 stream recomputation and one-source workspace."""
+    """Fused VJP with FP32 scalar recomputation and one-source workspace."""
     b, t, d = sources[0].shape
     n, m, ksum = b*t, len(widths)+1, sum(widths)
+    if not n:
+        return ([torch.empty_like(source) for source in sources],
+                [torch.zeros_like(bank) for bank in banks], torch.zeros_like(query))
     sources = [x.contiguous() for x in sources]
     banks, query = [x.contiguous() for x in banks], query.contiguous()
     valid, position, grad = valid.contiguous(), position.contiguous(), grad.contiguous()
-    # Recompute FP32 retrieved streams instead of saving them for every layer.
-    z = torch.empty((b,t,m,d),device=grad.device,dtype=torch.float32)
-    lse_recomputed = torch.empty_like(lse)
-    _msar_read_kernel[(n,m)](tuple(sources), tuple(banks), query, valid, position, z, lse_recomputed,
-        t,d,len(sources),m,tuple(widths),eps,
-        triton.next_power_of_2(d),num_warps=4)
-    # Only one source's view gradients are materialized at a time.
-    dv = torch.empty((m,n,d),device=grad.device,dtype=torch.float32)
+    single = len(sources) == 1
+    # ds_i = p_i * (<g,v_i> - <g,z>); recompute the scalar <g,z>
+    # directly rather than materializing all FP32 stream vectors z. A single
+    # finite source has ds=0, so that specialization skips recomputation.
+    gdot = lse
+    if not single:
+        gdot = torch.empty((b,t,m),device=grad.device,dtype=torch.float32)
+        _msar_read_kernel[(n,m)](tuple(sources), tuple(banks), query, valid, position, gdot, lse, grad,
+            t,d,len(sources),m,tuple(widths),eps,
+            triton.next_power_of_2(d),True,num_warps=4)
+    # Process complete sequences together so convolution never crosses a tile.
+    # Bound temporary storage automatically by both geometry and device memory;
+    # at least one sequence fits even when a single sequence exceeds the budget.
+    scratch_limit = min(128 * 1024**2, torch.cuda.get_device_properties(grad.device).total_memory // 64)
+    bytes_per_sequence = max(1, t * m * d * 4)
+    batch_tile = max(1, min(b, scratch_limit // bytes_per_sequence))
+    max_rows = batch_tile * t
+    # One source at a time, and only its convolution views: the identity VJP
+    # is computed in the input-transpose kernel instead of stored in FP32.
+    dv = torch.empty(((m-1)*max_rows*d,),device=grad.device,dtype=torch.float32)
     qrows = max(1, min(8, 8192 // triton.next_power_of_2(d)))
-    qgroups = triton.cdiv(n,qrows)
-    qpart = torch.empty((m,qgroups,d),device=grad.device,dtype=torch.float32)
-    wgroups = triton.cdiv(n,64)
-    wpart = torch.empty((ksum,wgroups,d),device=grad.device,dtype=torch.float32)
+    max_qgroups = triton.cdiv(max_rows,qrows)
+    # A single source has the same zero/NaN query partial in every dimension;
+    # the general multi-source path still needs one partial per dimension.
+    qpart_width = 1 if single else d
+    qpart = torch.empty((m*max_qgroups*qpart_width,),device=grad.device,dtype=torch.float32)
+    max_wgroups = triton.cdiv(max_rows,64)
+    wpart = torch.empty((ksum*max_wgroups*d,),device=grad.device,dtype=torch.float32)
     dq = torch.zeros_like(query, dtype=torch.float32)
     dws = []
     dxs = []
     for source, bank in zip(sources, banks):
-        _msar_vjp_value[(qgroups,m)](source,bank,query,valid,position,lse,z,grad,dv,qpart,
-            n,t,d,m,tuple(widths),eps,qrows,triton.next_power_of_2(d),qgroups,num_warps=4)
         dx = torch.empty_like(source)
-        _msar_vjp_source[(triton.cdiv(n*d,256),)](dv,bank,valid,position,dx,n,t,d,
-            tuple(widths),256,num_warps=4)
-        _msar_vjp_bank[(wgroups,triton.cdiv(d,32))](source,dv,valid,position,wpart,
-            n,t,d,tuple(widths),64,32,wgroups,num_warps=4)
-        dq.add_(qpart.sum(dim=1))
-        dws.append(wpart.sum(dim=1).to(bank.dtype))
+        dw = torch.empty_like(bank, dtype=torch.float32)
+        for start in range(0,b,batch_tile):
+            end = min(start+batch_tile,b)
+            rows = (end-start)*t
+            qgroups = triton.cdiv(rows,qrows)
+            wgroups = triton.cdiv(rows,64)
+            _msar_vjp_value[(qgroups,m-1)](source[start:end],bank,query,valid[start:end],
+                position[start:end],lse[start:end],gdot[start:end],grad[start:end],dv,qpart,
+                rows,t,d,m,tuple(widths),eps,qrows,triton.next_power_of_2(d),qgroups,single,num_warps=4)
+            _msar_vjp_source[(qgroups,)](source[start:end],bank,query,valid[start:end],
+                position[start:end],lse[start:end],gdot[start:end],grad[start:end],dv,qpart,dx[start:end],
+                rows,t,d,m,tuple(widths),eps,qrows,triton.next_power_of_2(d),qgroups,single,num_warps=4)
+            _msar_vjp_bank[(wgroups,triton.cdiv(d,32))](source[start:end],dv,valid[start:end],
+                position[start:end],wpart,rows,t,d,tuple(widths),64,32,wgroups,num_warps=4)
+            _msar_reduce_query_bank[((m+ksum)*triton.cdiv(d,32),)](
+                qpart,wpart,dq,dw,d,m,ksum,qgroups,wgroups,32,128,
+                start == 0,single,num_warps=4)
+        dws.append(dw.to(bank.dtype))
         dxs.append(dx)
     return dxs, dws, dq.to(query.dtype).contiguous()
 
@@ -7032,6 +7175,335 @@ def _msar_start_branch(sources, partial, branch, block_size):
     return sources, partial
 
 
+# Linear Depth Memory (RLS), adapted from the authors' flexible backend,
+# github.com/davidgonmar/depth-memory, commit 9a6865feeb43a153a4501cc9a6457e047ffef3c3.
+# The live memory is unwound in backward, so BF16 restoration is approximate.
+# The symmetric covariance tape stores only its upper triangle as in paper Appendix F.
+# MIT License
+# Copyright (c) 2026 David Gonzalez Martinez
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+if triton is not None:
+    @triton.jit
+    def _ldm_forward(M, P, COV, KEY, VALUE, QUERY, LAM, OUT, ERROR, GAIN, DEN,
+                     N: tl.constexpr, K: tl.constexpr, D: tl.constexpr,
+                     BK: tl.constexpr, BD: tl.constexpr, BN: tl.constexpr):
+        """Write/read RLS and save the symmetric covariance as a triangle."""
+        row = tl.program_id(0)*BN + tl.arange(0, BN)
+        kk = tl.arange(0, BK)
+        dd = tl.arange(0, BD)
+        row_ok = row < N
+        km = row_ok[:, None] & (kk[None, :] < K)
+        vm = row_ok[:, None] & (dd[None, :] < D)
+        key = tl.load(KEY + row[:, None]*K + kk[None, :], km, 0).to(tl.float32)
+        value = tl.load(VALUE + row[:, None]*D + dd[None, :], vm, 0).to(tl.float32)
+        lam = tl.load(LAM + row, row_ok, 1.).to(tl.float32)
+        u = tl.zeros((BN, BK), tl.float32)
+        den = lam
+        for i in tl.static_range(K):
+            p_off = row[:, None]*K*K + i*K + kk[None, :]
+            p_row = tl.load(P + p_off, km, 0).to(tl.float32)
+            # The live covariance starts at I and the RLS update preserves exact
+            # symmetry. Store its upper triangle only; backward mirrors indices.
+            packed = row[:, None] * (K*(K+1)//2) + i*(2*K-i+1)//2 + kk[None, :] - i
+            tl.store(COV + packed, p_row, km & (kk[None, :] >= i))
+            ui = tl.sum(p_row*key, 1)
+            ki = tl.load(KEY + row*K + i, row_ok, 0).to(tl.float32)
+            den += ki*ui
+            u = tl.where(kk[None, :] == i, ui[:, None], u)
+        # The tape is produced by the full CTA; synchronize before reading it.
+        tl.debug_barrier()
+        for i in tl.static_range(K):
+            p_off = row[:, None]*K*K + i*K + kk[None, :]
+            packed = row[:, None] * (K*(K+1)//2) + i*(2*K-i+1)//2 + kk[None, :] - i
+            upper = km & (kk[None, :] >= i)
+            p_row = tl.load(COV + packed, upper, 0).to(tl.float32)
+            ui = tl.sum(tl.where(kk[None, :] == i, u, 0.), 1)
+            p_new = (p_row - ui[:, None]*u/den[:, None])/lam[:, None]
+            tl.store(P + p_off, p_new, upper)
+            transpose_off = row[:, None]*K*K + kk[None, :]*K + i
+            tl.store(P + transpose_off, p_new, km & (kk[None, :] > i))
+        error = value
+        for i in tl.static_range(K):
+            m_off = row[:, None]*K*D + i*D + dd[None, :]
+            mr = tl.load(M + m_off, vm, 0).to(tl.float32)
+            ki = tl.load(KEY + row*K + i, row_ok, 0).to(tl.float32)
+            error -= ki[:, None]*mr
+        output = tl.zeros((BN, BD), tl.float32)
+        for i in tl.static_range(K):
+            m_off = row[:, None]*K*D + i*D + dd[None, :]
+            old = tl.load(M + m_off, vm, 0).to(tl.float32)
+            ui = tl.sum(tl.where(kk[None, :] == i, u, 0.), 1)
+            gain = ui/den
+            after = (old + gain[:, None]*error).to(M.dtype.element_ty)
+            tl.store(M + m_off, after, vm)
+            qi = tl.load(QUERY + row*K + i, row_ok, 0).to(tl.float32)
+            output += after.to(tl.float32)*qi[:, None]
+            tl.store(GAIN + row*K + i, gain, row_ok)
+        tl.store(OUT + row[:, None]*D + dd[None, :], output.to(OUT.dtype.element_ty), vm)
+        tl.store(ERROR + row[:, None]*D + dd[None, :], error.to(ERROR.dtype.element_ty), vm)
+        tl.store(DEN + row, den, row_ok)
+
+
+    @triton.jit
+    def _ldm_backward(M, P, COV, GM_IN, GP_IN, GM_OUT, GP_OUT, KEY, QUERY, LAM, ERROR, GAIN, DEN, DY,
+                      DK, DV, DQ, DL, N: tl.constexpr, K: tl.constexpr, D: tl.constexpr,
+                      BK: tl.constexpr, BD: tl.constexpr, BN: tl.constexpr, FIRST: tl.constexpr,
+                      HAS_GM: tl.constexpr, HAS_GP: tl.constexpr):
+        """Reverse RLS without mutating or pre-copying incoming cotangents."""
+        row = tl.program_id(0)*BN + tl.arange(0, BN)
+        ii = tl.arange(0, BK)
+        jj = tl.arange(0, BD)
+        row_ok = row < N
+        km = row_ok[:, None] & (ii[None, :] < K)
+        vm = row_ok[:, None] & (jj[None, :] < D)
+        mm = row_ok[:, None, None] & (ii[None, :, None] < K) & (jj[None, None, :] < D)
+        pm = row_ok[:, None, None] & (ii[None, :, None] < K) & (ii[None, None, :] < K)
+        mo = row[:, None, None]*K*D + ii[None, :, None]*D + jj[None, None, :]
+        po = row[:, None, None]*K*K + ii[None, :, None]*K + ii[None, None, :]
+        ko = row[:, None]*K + ii[None, :]
+        vo = row[:, None]*D + jj[None, :]
+        after = tl.load(M + mo, mm, 0).to(tl.float32)
+        p_after = tl.load(P + po, pm, 0).to(tl.float32)
+        lo = tl.minimum(ii[:, None], ii[None, :])
+        hi = tl.maximum(ii[:, None], ii[None, :])
+        packed = row[:, None, None]*(K*(K+1)//2) + lo[None, :, :]*(2*K-lo[None, :, :]+1)//2 + hi[None, :, :]-lo[None, :, :]
+        p_before = tl.load(COV + packed, pm, 0).to(tl.float32)
+        key = tl.load(KEY + ko, km, 0).to(tl.float32)
+        query = tl.load(QUERY + ko, km, 0).to(tl.float32)
+        error = tl.load(ERROR + vo, vm, 0).to(tl.float32)
+        gain = tl.load(GAIN + ko, km, 0).to(tl.float32)
+        lam = tl.load(LAM + row, row_ok, 1.).to(tl.float32)
+        den = tl.load(DEN + row, row_ok, 1.).to(tl.float32)
+        dy = tl.load(DY + vo, vm, 0).to(tl.float32)
+        before = (after - gain[:, :, None]*error[:, None, :]).to(M.dtype.element_ty).to(tl.float32)
+        if FIRST:
+            before = tl.zeros((BN, BK, BD), tl.float32)
+            after = (gain[:, :, None]*error[:, None, :]).to(M.dtype.element_ty).to(tl.float32)
+        read_grad = (query[:, :, None]*dy[:, None, :]).to(M.dtype.element_ty).to(tl.float32)
+        if HAS_GM:
+            gm_in = tl.load(GM_IN + mo, mm, 0).to(tl.float32)
+        else:
+            gm_in = tl.zeros((BN, BK, BD), tl.float32)
+        h = (gm_in + read_grad).to(M.dtype.element_ty).to(tl.float32)
+        dq = tl.sum(after*dy[:, None, :], 2)
+        dv = tl.sum(h*gain[:, :, None], 1)
+        dg = tl.sum(h*error[:, None, :], 2)
+        if HAS_GP:
+            jp = tl.load(GP_IN + po, pm, 0).to(tl.float32)
+        else:
+            jp = tl.zeros((BN, BK, BK), tl.float32)
+        u = gain*den[:, None]
+        ju = tl.sum(jp*u[:, None, :], 2)
+        jtu = tl.sum(jp*u[:, :, None], 1)
+        du = dg/den[:, None] - (ju+jtu)/(lam*den)[:, None]
+        dden = -tl.sum(dg*gain, 1)/den + tl.sum(u*ju, 1)/(lam*den*den)
+        du += dden[:, None]*key
+        dk = -tl.sum(before*dv[:, None, :], 2) + dden[:, None]*u + tl.sum(p_before*du[:, :, None], 1)
+        gp = jp/lam[:, None, None] + du[:, :, None]*key[:, None, :]
+        dl = dden - tl.sum(tl.sum(jp*p_after, 2), 1)/lam
+        gm = h - key[:, :, None]*dv[:, None, :]
+        tl.store(GM_OUT + mo, gm.to(GM_OUT.dtype.element_ty), mm)
+        tl.store(GP_OUT + po, gp, pm)
+        tl.store(DK + ko, dk.to(DK.dtype.element_ty), km)
+        tl.store(DQ + ko, dq.to(DQ.dtype.element_ty), km)
+        tl.store(DV + vo, dv.to(DV.dtype.element_ty), vm)
+        tl.store(DL + row, dl, row_ok)
+        tl.store(M + mo, before.to(M.dtype.element_ty), mm)
+        tl.store(P + po, p_before, pm)
+
+
+
+def _ldm_kernel_shape(memory):
+    k, d = memory.shape[-2:]
+    return memory.numel()//(k*d), k, d, max(2, triton.next_power_of_2(k)), max(2, triton.next_power_of_2(d))
+
+
+def _ldm_row_tile(bk, bd, budget, maximum):
+    rows = max(1, min(maximum, budget//(bk*bd)))
+    return 1 << (rows.bit_length()-1)
+
+
+@torch.library.custom_op("neollm_ldm::forward", mutates_args=("memory", "covariance"))
+def _ldm_forward_op(memory: torch.Tensor, covariance: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
+                query: torch.Tensor, forgetting: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    out, error = torch.empty_like(value), torch.empty_like(value)
+    gain = torch.empty_like(key, dtype=torch.float32)
+    den = torch.empty_like(forgetting, dtype=torch.float32)
+    k = memory.shape[-2]
+    before_cov = torch.empty((*covariance.shape[:-2], k*(k+1)//2), dtype=covariance.dtype, device=covariance.device)
+    n, k, d, bk, bd = _ldm_kernel_shape(memory)
+    if n:
+        bn = _ldm_row_tile(bk, bd, 8192, 16)
+        _ldm_forward[(triton.cdiv(n, bn),)](memory, covariance, before_cov, key, value, query, forgetting,
+            out, error, gain, den, n, k, d, bk, bd, bn, num_warps=4)
+    return out, error, gain, den, before_cov
+
+
+@_ldm_forward_op.register_fake
+def _ldm_forward_fake(memory, covariance, key, value, query, forgetting):
+    k = memory.shape[-2]
+    return (torch.empty_like(value), torch.empty_like(value), torch.empty_like(key, dtype=torch.float32),
+            torch.empty_like(forgetting, dtype=torch.float32),
+            torch.empty((*covariance.shape[:-2], k*(k+1)//2), dtype=covariance.dtype, device=covariance.device))
+
+
+@torch.library.custom_op("neollm_ldm::backward", mutates_args=("memory", "covariance"))
+def _ldm_backward_op(memory: torch.Tensor, covariance: torch.Tensor, before_cov: torch.Tensor,
+                 grad_memory: torch.Tensor | None, grad_covariance: torch.Tensor | None, key: torch.Tensor,
+                 query: torch.Tensor, forgetting: torch.Tensor, error: torch.Tensor,
+                 gain: torch.Tensor, den: torch.Tensor, grad_out: torch.Tensor,
+                 first: bool) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    dk, dv, dq, dl = (torch.empty_like(item) for item in (key, error, query, forgetting))
+    gm, gp = torch.empty_like(memory), torch.empty_like(covariance)
+    gm_in = grad_memory.contiguous() if grad_memory is not None else memory
+    gp_in = grad_covariance.contiguous() if grad_covariance is not None else covariance
+    n, k, d, bk, bd = _ldm_kernel_shape(memory)
+    if n:
+        bn = _ldm_row_tile(bk, bd, 2048, 8)
+        _ldm_backward[(triton.cdiv(n, bn),)](memory, covariance, before_cov, gm_in, gp_in, gm, gp,
+            key, query, forgetting, error, gain, den, grad_out, dk, dv, dq, dl,
+            n, k, d, bk, bd, bn, first, grad_memory is not None, grad_covariance is not None, num_warps=4)
+    return gm, gp, dk, dv, dq, dl
+
+
+@_ldm_backward_op.register_fake
+def _ldm_backward_fake(memory, covariance, before_cov, grad_memory, grad_covariance, key, query,
+                   forgetting, error, gain, den, grad_out, first):
+    return tuple(torch.empty_like(item) for item in (memory, covariance, key, error, query, forgetting))
+
+
+class _LDMReversibleStep(torch.autograd.Function):
+    """First-order reverse-mode update with an explicitly owned mutable state."""
+
+    @staticmethod
+    def forward(ctx, memory, covariance, key, value, query, forgetting, first):
+        out, error, gain, den, before_cov = _ldm_forward_op(
+            memory, covariance, key.contiguous(), value.contiguous(), query.contiguous(), forgetting.contiguous())
+        ctx.mark_dirty(memory, covariance)
+        ctx.set_materialize_grads(False)
+        ctx.memory, ctx.covariance = memory.detach(), covariance.detach()
+        ctx.first, ctx.used = first, False
+        ctx.save_for_backward(before_cov, key, query, forgetting, error, gain, den)
+        return memory, covariance, out
+
+    @staticmethod
+    def backward(ctx, grad_memory, grad_covariance, grad_out):
+        if ctx.used:
+            raise RuntimeError("LDM reversible state supports one backward per forward")
+        ctx.used = True
+        before_cov, key, query, forgetting, error, gain, den = ctx.saved_tensors
+        dy = torch.zeros_like(error) if grad_out is None else grad_out.contiguous()
+        gm, gp, dk, dv, dq, dl = _ldm_backward_op(ctx.memory, ctx.covariance, before_cov,
+                                              grad_memory, grad_covariance, key.contiguous(),
+                                              query.contiguous(), forgetting.contiguous(),
+                                              error, gain, den, dy, ctx.first)
+        return gm, gp, dk, dv, dq, dl, None
+
+
+def _ldm_step(memory, covariance, key, value, query, forgetting, first):
+    return _LDMReversibleStep.apply(memory, covariance, key, value, query, forgetting, first)
+
+
+def _ldm_reference_step(memory, covariance, key, value, query, forgetting):
+    """Non-mutating RLS reference, also used where Triton is unavailable."""
+    k, v, q = key.float(), value.float(), query.float()
+    u = torch.einsum("...ij,...j->...i", covariance, k)
+    den = forgetting + (k * u).sum(-1)
+    gain = u / den.unsqueeze(-1)
+    error = v - torch.einsum("...kd,...k->...d", memory.float(), k)
+    updated = (memory.float() + gain.unsqueeze(-1) * error.unsqueeze(-2)).to(memory.dtype)
+    covariance_new = (covariance - u.unsqueeze(-1) * u.unsqueeze(-2)
+                      / den[..., None, None]) / forgetting[..., None, None]
+    read = torch.einsum("...kd,...k->...d", updated.float(), q).to(value.dtype)
+    return updated, covariance_new, read
+
+
+class LDMAddress(nn.Module):
+    """Small input-dependent address and identity-initialized group gate."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.groups = config.ldm_groups
+        self.key_dim = config.ldm_key_dim
+        self.context_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.w1 = nn.Linear(config.hidden_size, config.ldm_hidden_dim, bias=True)
+        self.w2 = nn.Linear(config.ldm_hidden_dim, self.groups * self.key_dim, bias=True)
+        self.gate = nn.Linear(config.ldm_hidden_dim, self.groups, bias=True)
+        # NeoLLM's HF post_init visits child Linear modules independently.
+        self.gate._ldm_gate_zero_init = True
+        nn.init.zeros_(self.gate.weight)
+        nn.init.zeros_(self.gate.bias)
+
+    def forward(self, source):
+        normalized = F.rms_norm(source.float(), self.context_norm.normalized_shape,
+                                self.context_norm.weight.float(), self.context_norm.eps)
+        features = F.silu(self.w1(normalized.to(source.dtype)))
+        address = self.w2(features).float().reshape(*source.shape[:-1], self.groups, self.key_dim)
+        address = F.normalize(address, p=2, dim=-1, eps=1e-6)
+        gate = 2.0 * torch.sigmoid(self.gate(features).float())
+        return address, gate
+
+
+class LDMWrite(nn.Module):
+    """Generate a weighted write from a residual context and branch output."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.address = LDMAddress(config)
+        self.groups = config.ldm_groups
+        self.value_dim = config.hidden_size // self.groups
+        self.forgetting_logit = nn.Parameter(torch.full((self.groups,), math.log(.99 / .01)))
+
+    def forward(self, context, value, valid, dtype):
+        key, weight = self.address(context)
+        # Mask after sqrt: sqrt(0) in an autograd path can produce NaN gradients.
+        # Match the authors' numerical floors for an extreme negative gate
+        # or forgetting logit, avoiding sqrt(0) and division by zero.
+        scale = weight.clamp_min(1e-6).sqrt() * valid.unsqueeze(-1).to(weight.dtype)
+        key = (key * scale.unsqueeze(-1)).to(dtype).contiguous()
+        value = (value.float().reshape(*value.shape[:-1], self.groups, self.value_dim)
+                 * scale.unsqueeze(-1)).to(dtype).contiguous()
+        forgetting = torch.sigmoid(self.forgetting_logit.float()).clamp_min(1e-6).expand_as(weight)
+        forgetting = torch.where(valid.unsqueeze(-1), forgetting, torch.ones_like(forgetting)).contiguous()
+        return key, value, forgetting
+
+
+class LDMRead(nn.Module):
+    """Read the latest pending write with a normalized gated query."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.address = LDMAddress(config)
+        self.read_gain = nn.Parameter(torch.ones(config.ldm_groups))
+
+    def forward(self, context, state, valid):
+        memory, covariance, key, value, forgetting, first = state
+        query, gate = self.address(context)
+        query = (query * (gate * self.read_gain.float()).unsqueeze(-1)
+                 * valid[..., None, None].to(query.dtype)).to(memory.dtype).contiguous()
+        if memory.is_cuda and triton is not None:
+            memory, covariance, read = _ldm_step(memory, covariance, key, value, query, forgetting, first)
+        else:
+            memory, covariance, read = _ldm_reference_step(memory, covariance, key, value, query, forgetting)
+        return memory, covariance, read.reshape_as(context).to(context.dtype)
+
+
 class NeoLLMDecoderLayer(GradientCheckpointingLayer):
     """
     Decoder layer with standard residual connections.
@@ -7205,6 +7677,11 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         self.msar_block_branches = config.msar_block_branches
         self.msar_attn_read = MSARRead(config) if self.use_msar else None
         self.msar_mlp_read = MSARRead(config) if self.use_msar else None
+        self.use_ldm = bool(getattr(config, "use_ldm", False))
+        self.ldm_attn_read = LDMRead(config) if self.use_ldm else None
+        self.ldm_attn_write = LDMWrite(config) if self.use_ldm else None
+        self.ldm_mlp_read = LDMRead(config) if self.use_ldm else None
+        self.ldm_mlp_write = LDMWrite(config) if self.use_ldm else None
         self.use_attn_res = getattr(config, "use_attn_res", False)
         if self.use_attn_res:
             self.attn_res_query_attn = nn.Parameter(torch.zeros(config.hidden_size))
@@ -7500,19 +7977,43 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         dag_sources_y=None,
         dag_plan=None,
         dynamics_valid_mask: Optional[torch.Tensor] = None,
+        msar_sources=None,
+        msar_partial=None,
+        msar_banks=None,
+        msar_valid=None,
+        msar_position=None,
+        ldm_state=None,
+        ldm_valid=None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple:
-        """Update coupled X/Y streams with optional global/local DAG corrections.
-
-        Raw X/Y histories retain separate residual reads and shared target-layer
-        attention reads. Existing stream scaling, GPAS and memory policies remain
-        in their original order. A missing dag_plan selects the legacy path.
+        """Update X/Y streams with the selected optional depth mechanism.
+
+        In the hybrid, MSAR stores one history of unscaled shared branch outputs.
+        Its causal depth read is added after separate X/Y normalization, before
+        their optional shared input norm. This uses the existing SiameseNorm
+        input norm and does not introduce an additional normalization. The X/Y
+        update equations and final merge are unchanged, while branch values can
+        change.
+        This hybrid is an adaptation, not the single-stream paper architecture.
+
+        LDM likewise stores one shared history of unscaled branch outputs. It
+        adds its read to the separate X/Y normalization result before the
+        existing shared input norm; write keys use the mean of the updated X/Y
+        streams. Its kernel is independent of the SiameseNorm kernel.
         """
         # SiameseNorm keeps two coupled streams with shared Attention/MLP
         # 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(
+                msar_sources, msar_partial, branch, self.msar_block_branches)
+            msar_attn = self.msar_attn_read(msar_sources, msar_partial, msar_banks,
+                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)
@@ -7520,17 +8021,23 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
                 dtype=x_attn_norm.dtype, device=x_attn_norm.device
             )
             h_attn = x_scale * x_attn_norm + y_attn_norm
-            if self.siamese_attn_input_norm is not None:
-                h_attn = self.siamese_attn_input_norm(h_attn)
         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 self.siamese_attn_input_norm is not None else None,
+                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,
             )
+        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):
+            h_attn = self.siamese_attn_input_norm(h_attn)
 
         stack_metrics = None
         if self.use_stack_memory:
@@ -7566,36 +8073,56 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             y_states = y_states + dag_ry
         x_after_attn = x_states + stream_scale * attn_out
         y_after_attn = y_states + attn_out
+        if self.use_ldm:
+            ldm_state = (ldm_memory, ldm_covariance,
+                         *self.ldm_attn_write((x_after_attn + y_after_attn) * .5,
+                                              attn_out, ldm_valid, ldm_memory.dtype), False)
+        if self.use_msar:
+            msar_partial = attn_out if msar_partial is None else msar_partial + attn_out
 
         if self.use_gpas:
             x_after_attn = self.gpas_attn(x_after_attn)
 
         # ── MLP shared block ──────────────────────────────────────────────
+        if self.use_msar:
+            branch = 2 * self.layer_idx + 1
+            msar_sources, msar_partial = _msar_start_branch(
+                msar_sources, msar_partial, branch, self.msar_block_branches)
+            msar_mlp = self.msar_mlp_read(msar_sources, msar_partial, msar_banks,
+                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
-            if self.siamese_mlp_input_norm is not None:
-                h_mlp = self.siamese_mlp_input_norm(h_mlp)
         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 self.siamese_mlp_input_norm is not None else None,
+                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 else None
+                    if self.use_lns and not (self.use_msar or self.use_ldm) else None
                 ),
                 out_dtype=(
                     torch.get_autocast_dtype("cuda")
-                    if torch.is_autocast_enabled("cuda") else None
+                    if not (self.use_msar or self.use_ldm) and torch.is_autocast_enabled("cuda") else None
                 ),
             )
+        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):
+            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" else h_mlp
+        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
         delta_m = self.mlp(h_lns2)
 
         jtok_aux_stats = None
@@ -7616,11 +8143,21 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
 
         x_next = x_after_attn + stream_scale * shared_update
         y_next = y_after_attn + shared_update
+        if self.use_ldm:
+            ldm_state = (ldm_memory, ldm_covariance,
+                         *self.ldm_mlp_write((x_next + y_next) * .5,
+                                             shared_update, ldm_valid, ldm_memory.dtype), False)
+        if self.use_msar:
+            msar_partial = shared_update if msar_partial is None else msar_partial + shared_update
 
         if self.use_gpas:
             x_next = self.gpas_mlp(x_next)
 
         outputs = (x_next, y_next)
+        if self.use_ldm:
+            outputs += (ldm_state,)
+        if self.use_msar:
+            outputs += (msar_sources, msar_partial)
         if self.use_stack_memory:
             outputs += (stack_memory, stack_memory_mask, stack_metrics)
         if jtok_aux_stats is not None:
@@ -7642,6 +8179,8 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         msar_banks=None,
         msar_valid=None,
         msar_position=None,
+        ldm_state=None,
+        ldm_valid=None,
         output_attentions: Optional[bool] = False,
         repo_rope_args: Optional[Tuple[torch.Tensor, float]] = None,
         position_ids: Optional[torch.LongTensor] = None,
@@ -7690,6 +8229,10 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             h_attn = hidden_states
             residual_attn = hidden_states
 
+        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
+
         # ── Attention block ───────────────────────────────────────────────
 
         # Optional embedding input normalisation:
@@ -7740,6 +8283,9 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             attn_aug = residual_attn + hidden_states
 
         h_tilde = self.gpas_attn(attn_aug) if self.use_gpas else attn_aug
+        if self.use_ldm:
+            ldm_state = (ldm_memory, ldm_covariance,
+                         *self.ldm_attn_write(h_tilde, hidden_states, ldm_valid, ldm_memory.dtype), False)
 
         # ── Attention Residuals: compute pre-MLP input ────────────────────
         # After attention, the partial sum is updated with h_tilde.
@@ -7764,6 +8310,10 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             h_mlp = h_tilde
             residual_mlp = h_tilde
 
+        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
+
         # ── MLP block ─────────────────────────────────────────────────────
         h_normed2 = _apply_norm(self.post_attention_layernorm, h_mlp)
         h_lns2 = self.lns_mlp(h_normed2) if self.use_lns else h_normed2
@@ -7802,8 +8352,13 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             mlp_aug = residual_mlp + delta_m
 
         hidden_states = self.gpas_mlp(mlp_aug) if self.use_gpas else mlp_aug
+        if self.use_ldm:
+            ldm_state = (ldm_memory, ldm_covariance,
+                         *self.ldm_mlp_write(hidden_states, delta_m, ldm_valid, ldm_memory.dtype), False)
 
         outputs = (hidden_states,)
+        if self.use_ldm:
+            outputs += (ldm_state,)
         if self.use_msar:
             outputs += (msar_sources, msar_partial)
         if jtok_aux_stats is not None:
@@ -8147,6 +8702,9 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
     _supports_flash_attn_2 = True
     _supports_sdpa = True
     _is_stateful = True
+    # Preserve the paper's FP32 LDM controls while loading a BF16 checkpoint.
+    # Transformers applies this before copying tensors, avoiding lossy casts.
+    _keep_in_fp32_modules_strict = {"ldm_*.read_gain", "ldm_*.forgetting_logit"}
 
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
@@ -8250,9 +8808,9 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
         return model
 
     def _init_weights(self, module):
-        """Apply HF initialization, then enforce exact-zero marked DAG heads."""
+        """Apply HF initialization, then preserve identity-centered gates."""
         super()._init_weights(module)
-        if getattr(module, "_dag_zero_init", False):
+        if getattr(module, "_dag_zero_init", False) or getattr(module, "_ldm_gate_zero_init", False):
             nn.init.zeros_(module.weight)
             nn.init.zeros_(module.bias)
             return
@@ -8520,6 +9078,8 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 sum(config.msar_kernel_sizes), config.hidden_size))
             if config.use_msar else None
         )
+        self.ldm_embedding_write = LDMWrite(config) if config.use_ldm else None
+        self.ldm_final_read = LDMRead(config) if config.use_ldm else None
         self.dagformer = (DAGFormerTopology(config, self.layers)
                           if bool(getattr(config, "use_dagformer", False)) else None)
         self.use_siamesenorm = bool(getattr(config, "use_siamesenorm", False))
@@ -8862,6 +9422,26 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         hidden_states = inputs_embeds
         use_siamesenorm = bool(getattr(self.config, "use_siamesenorm", False))
         siamese_y_states = inputs_embeds if use_siamesenorm else None
+        use_ldm = bool(getattr(self.config, "use_ldm", False))
+        if use_ldm:
+            if self.gradient_checkpointing and self.training:
+                raise ValueError("LDM reversible state is incompatible with whole-layer gradient checkpointing.")
+            if attention_mask is not None and (attention_mask.ndim != 2 or attention_mask.shape != hidden_states.shape[:2]):
+                raise ValueError("LDM requires a token-domain 2-D attention mask.")
+            ldm_valid = (attention_mask.to(dtype=torch.bool).contiguous() if attention_mask is not None
+                         else torch.ones(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device))
+            state_dtype = (torch.get_autocast_dtype("cuda") if hidden_states.is_cuda and
+                           torch.is_autocast_enabled("cuda") else hidden_states.dtype)
+            groups, key_dim = self.config.ldm_groups, self.config.ldm_key_dim
+            value_dim = self.config.hidden_size // groups
+            ldm_memory = torch.zeros((*hidden_states.shape[:2], groups, key_dim, value_dim),
+                                     device=hidden_states.device, dtype=state_dtype)
+            ldm_covariance = torch.eye(key_dim, device=hidden_states.device, dtype=torch.float32).expand(
+                *hidden_states.shape[:2], groups, key_dim, key_dim).clone()
+            ldm_state = (ldm_memory, ldm_covariance,
+                         *self.ldm_embedding_write(hidden_states, hidden_states, ldm_valid, state_dtype), True)
+        else:
+            ldm_state = ldm_valid = None
         # The same validated dense/right-padding training contract applies to
         # the causal planner: valid queries cannot see future padded keys.
         dag_plans = (self.dagformer(input_ids, inputs_embeds, position_ids,
@@ -9021,18 +9601,28 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     dag_sources=dag_sources,
                     dag_plan=dag_plans[layer_idx] if dag_plans is not None else None,
                     dynamics_valid_mask=attention_mask,
+                    msar_sources=msar_sources, msar_partial=msar_partial,
+                    msar_banks=self.msar_banks, msar_valid=msar_valid,
+                    msar_position=msar_position,
+                    ldm_state=ldm_state, ldm_valid=ldm_valid,
                     **kwargs,
                 )
                 hidden_states = layer_outputs[0]
                 siamese_y_states = layer_outputs[1]
-                if use_stack_memory:
-                    stack_memory = layer_outputs[2]
-                    stack_memory_mask = layer_outputs[3]
-                    stack_layer_metrics = layer_outputs[4]
-                    all_stack_metrics = all_stack_metrics + (stack_layer_metrics,)
-                    extras_start = 5
+                if use_msar:
+                    msar_sources, msar_partial = layer_outputs[2:4]
+                    extras_start = 4
+                elif use_ldm:
+                    ldm_state = layer_outputs[2]
+                    extras_start = 3
                 else:
                     extras_start = 2
+                if use_stack_memory:
+                    stack_memory = layer_outputs[extras_start]
+                    stack_memory_mask = layer_outputs[extras_start + 1]
+                    stack_layer_metrics = layer_outputs[extras_start + 2]
+                    all_stack_metrics = all_stack_metrics + (stack_layer_metrics,)
+                    extras_start += 3
             else:
                 layer_outputs = decoder_layer(
                     hidden_states,
@@ -9044,6 +9634,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     msar_sources=msar_sources, msar_partial=msar_partial,
                     msar_banks=self.msar_banks, msar_valid=msar_valid,
                     msar_position=msar_position,
+                    ldm_state=ldm_state, ldm_valid=ldm_valid,
                     output_attentions=output_attentions,
                     repo_rope_args=repo_rope_args,
                     position_ids=position_ids,
@@ -9061,6 +9652,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     extras_start = 3
                     # The last branch's h_mlp + delta_m is the terminal
                     # residual merge. No new source expansion is applied.
+                elif use_ldm:
+                    ldm_state = layer_outputs[1]
+                    extras_start = 2
                 else:
                     extras_start = 1
 
@@ -9106,7 +9700,11 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
 
         if use_siamesenorm:
             norm_backend = getattr(self, "siamesenorm_kernel_backend", "reference")
-            if norm_backend == "reference":
+            if use_ldm:
+                merged = self.siamese_x_final_norm(hidden_states) + self.siamese_y_final_norm(siamese_y_states)
+                _, _, ldm_final = self.ldm_final_read(merged, ldm_state, ldm_valid)
+                hidden_states = self.siamese_final_norm(merged + ldm_final)
+            elif norm_backend == "reference":
                 hidden_states = self.siamese_final_norm(
                     self.siamese_x_final_norm(hidden_states)
                     + self.siamese_y_final_norm(siamese_y_states)
@@ -9120,6 +9718,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     eps=self.siamese_x_final_norm.eps, backend=norm_backend,
                 )
         else:
+            if use_ldm:
+                _, _, ldm_final = self.ldm_final_read(hidden_states, ldm_state, ldm_valid)
+                hidden_states = hidden_states + ldm_final
             hidden_states = self.norm(hidden_states)
 
         if output_hidden_states:
@@ -11431,6 +12032,17 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 for name, parameter in module.named_parameters():
                     if parameter.requires_grad:
                         groups[f"{kind}/layer_{module.layer_idx}/{name}"] = (parameter,)
+            for slot in ("attn", "mlp"):
+                read = getattr(layer, f"msar_{slot}_read", None)
+                if read is not None:
+                    prefix = f"msar/layer_{layer.layer_idx}/{slot}"
+                    if read.query.requires_grad:
+                        groups[f"{prefix}/query"] = (read.query,)
+                    mixer = tuple(parameter for parameter in
+                                  (read.mixer_weight, read.mixer_bias, read.mixer_scale)
+                                  if parameter.requires_grad)
+                    if mixer:
+                        groups[f"{prefix}/mixer"] = mixer
             families = {
                 "dag": ("dagformer_local",),
                 "attnres": ("attn_res_query_attn", "attn_res_query_mlp", "attn_res_norm"),
@@ -11453,6 +12065,8 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 groups[f"fan/layer_{layer.layer_idx}"] = (layer.self_attn.lambda_1, layer.self_attn.lambda_2)
         if self.model.dagformer is not None:
             groups["dag/global_topology"] = tuple(p for p in self.model.dagformer.parameters() if p.requires_grad)
+        if self.model.msar_banks is not None and self.model.msar_banks.requires_grad:
+            groups["msar/banks"] = (self.model.msar_banks,)
         return groups
 
     def set_input_embeddings(self, value):

optimizer_precision.json

@@ -8,9 +8,9 @@
   "delta_enabled": false,
   "delta_eta": 0.37,
   "fp32_master_active": true,
-  "fp32_master_elements": 85093304,
+  "fp32_master_elements": 84567992,
   "fp32_master_groups": 6,
-  "fp32_master_tensors": 378,
+  "fp32_master_tensors": 374,
   "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": 127,
+  "ordinary_linear_layers": 124,
   "parameter_elements_by_dtype": {
-    "bfloat16": 85091128,
+    "bfloat16": 84565816,
     "float32": 2176
   },
-  "parameter_storage_bytes": 170190960,
+  "parameter_storage_bytes": 169140336,
   "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:bc8f43459b8eea473f97704f0e1395ab4f28e9644cd98156e2a840b62cd44b23
+oid sha256:f25ab9f8e336f20fbb462d246d9a32edde7efc0113368fb5f2ca06f3cfe35459
 size 5329