Atomic-GDN2

Atomic-GDN2 is an experimental language model architecture combining Gated DeltaNet-2 (GDN-2), Multi-head Latent Attention (MLA), and hierarchical Block-Delta Attention Residuals (DAR).

This repository contains two checkpoints:

• Baseline: 28 GDN-2 layers.
• Hybrid: 21 GDN-2 layers, 7 MLA layers, and hierarchical DAR residual mixing.

Both models were trained from scratch on OpenWebText-gpt2 using a TPU v5e-8. The hybrid completed approximately 4.50B training tokens in 9.25 hours. The baseline was interrupted after approximately 3.71B tokens and 7.3 hours.

The primary comparison is validation bits per byte (BPB) at matched training-token budgets. Final checkpoint results are reported separately because the two runs processed different numbers of tokens.

1. Model summary

Parameter	Baseline	Hybrid
Parameters	334,289,920	342,189,056
Layers	28	28
GDN-2 layers	28	21
MLA layers	0	7
DAR	No	Yes
Hidden dimension	1024	1024
Attention heads	8	8
Head dimension	128	128
FFN hidden dimension	3584	3584
KV latent rank	—	512
QK non-RoPE dimension	—	64
QK RoPE dimension	—	64
DAR latent dimension	—	256
Vocabulary size	50,257	50,257
Context length	2048	2048
Activation	SwiGLU	SwiGLU
Dropout	0.1	0.1
Tied embeddings	Yes	Yes
Compute precision	bf16	bf16
Parameter precision	fp32	fp32

The hybrid contains approximately 2.36% more parameters than the baseline.

2. Architecture

2.1 Gated DeltaNet-2

GDN-2 uses a recurrent matrix state to implement gated linear attention.

S_t = (I − β_t k_t k_t^T) Diag(α_t) S_{t−1} + β_t k_t v_t^T
o_t = S_t^T q_t

Here, S_t is the recurrent state, β_t is a per-token update gate, and α_t contains per-channel decay factors.

The log-decay is parameterized as:

g_raw = W_g x + b_g
g     = −min(softplus(g_raw), 160 / bs2)
α     = exp(g)

With bs2 = 32, the configured domain is g ∈ [−5, 0], corresponding to α ∈ [e^−5, 1].

The kernel receives g directly and applies the exponential internally. The implementation does not use Mamba-2's A_log or dt_bias parameters.

The decay bound is a constraint of the current BTL kernel implementation rather than a general requirement of GDN-2.

2.2 Multi-head Latent Attention

The hybrid uses MLA in seven layers. Key-value representations are compressed into a lower-dimensional latent space and projected back into per-head keys and values.

c_KV   = x W_KV^D
k_nope = c_KV W_K^U
v      = c_KV W_V^U
k_rope = RoPE(x W_K^R)

The configuration uses a KV latent rank of 512, eight heads, and separate 64-dimensional non-RoPE and RoPE components.

Rotary positional embeddings are applied to the positional component separately.

MLA uses causal flash attention with query and key blocks of 512 tokens and an attention scale of 1 / sqrt(d_head).

2.3 Hierarchical Block-Delta Attention Residuals

DAR replaces conventional additive residual connections with learned mixing over sublayer deltas and block-level summaries.

A conventional residual connection is:

h_l = h_{l−1} + f_l(h_{l−1})

DAR instead constructs a source set containing the block input and the deltas produced by preceding sublayers.

For intra-block mixing:

S_i = [h_block, Δ_1, …, Δ_i]
a_i = softmax(q_i K_i^T / sqrt(d_latent))
h_i = Σ_j a_ij S_i,j

Cross-block mixing combines previously computed block summaries with local sublayer deltas.

Each block contains four sublayers. Its summary is formed by accumulating their deltas:

s_j = Σ_{i ∈ j} Δ_i

The configured DAR latent dimension is 256.

Before accumulation, DAR sublayer outputs pass through jnp.clip(x, -1e3, 1e3) and jnp.nan_to_num. Clipping is intended to prevent numerical blow-up during early training when delta magnitudes may be large. The clip bound of 1e3 is a fixed implementation choice and has not been tuned. The effect of this numerical safeguard on optimization has not been independently measured.

3. Training setup

Setting	Baseline	Hybrid
Hardware	TPU v5e-8	TPU v5e-8
Number of chips	8	8
Batch size	48	32
Sequence length	2048	2048
Steps per epoch	15,251	22,877
Planned epochs	3	3
Actual final step	37,700	68,631
Actual training tokens	~3.71B	4.497B
Wall-clock duration	~7.3 h	~9.25 h
TPU-hours	~58.4	~74.0
Run status	Interrupted	Completed

Token counts are calculated from batch size, sequence length, and optimizer steps:

T = steps × batch_size × sequence_length

The token budget was matched by design for the planned runs, but the baseline did not complete its full budget. Consequently, the actual training runs are not matched in total tokens.

The baseline was interrupted at approximately 3.71B tokens. Its final checkpoint is therefore an incomplete-run checkpoint, not a three-epoch result.

The hybrid completed approximately 4.50B tokens. Its longer wall-clock duration reflects the additional computation and smaller batch size.

3.1 Optimizer

Both models use Muon for eligible two-dimensional parameters and AdamW for the remaining parameters.

The parameter split is based on ndim == 2, excluding embeddings from Muon.

Hyperparameter	Value
Peak learning rate	3e-4
Schedule	Cosine decay
Warmup	1%
Weight decay	0.1
Global gradient clipping	1.0
AdamW beta1	0.9
AdamW beta2	0.95
Muon Newton–Schulz iterations	5

Muon uses the configured Newton–Schulz coefficients a = 3.4445, b = −4.7750, and c = 2.0315, with shape-dependent scaling.

AdamW is applied to one-dimensional parameters, biases, normalization parameters, and embeddings.

3.2 Precision

Model computation uses bf16 with fp32 parameters.

The GDN-2 kernel uses a bf16x3 configuration for supported operations, approximating fp32 arithmetic through multiple bf16 passes. This is not numerically identical to native fp32 execution.

3.3 Distributed execution

Training uses data parallelism with replicated parameters. The batch axis is sharded across eight TPU chips, and the Pallas kernels use shard_map for batch-axis distribution.

The baseline processes six sequences per chip per step, compared with four for the hybrid.

4. Results

4.1 Matched-token comparison

The primary comparison is at 1.5B training tokens, where both models have reported evaluation points.

Metric	Baseline	Hybrid
Tokens seen	1.5B	1.5B
Validation BPB	1.0949	1.0671

The hybrid's reported advantage at this matched-token checkpoint is:

Δ_BPB = 1.0949 − 1.0671 = 0.0278 ≈ 0.028

Lower BPB indicates better predictive compression of the evaluation text.

This is the primary comparison because both models are evaluated at the same training-token budget.

4.2 Later checkpoints

The following values refer to later checkpoints with different training budgets.

Metric	Baseline	Hybrid
Tokens seen	~3.71B	~4.50B
Best reported validation BPB	1.0361	1.0058
Final optimizer step	37,700	68,631
Wall-clock duration	~7.3 h	~9.25 h

The final optimizer step is not necessarily the step at which the best validation BPB was observed.

These values are reported separately from the matched-token comparison. The difference between them cannot be attributed exclusively to architecture because the models were trained for different numbers of tokens and optimizer steps.

The hybrid's reported BPB advantage is consistent with the matched-token result, but the later checkpoints do not constitute a compute-matched comparison.

4.3 Time-to-target

Time-to-target measures the number of training tokens required for each model to reach a specified validation BPB threshold.

Use the first observed evaluation point at or below each threshold, without interpolation between evaluation points.

Target BPB	Baseline tokens	Hybrid tokens	Token reduction
≤ 1.10	1.47B	0.98B	33%
≤ 1.08	2.36B	1.28B	46%
≤ 1.06	2.95B	2.36B	20%

Token reduction is calculated as:

1 − T_hybrid / T_baseline

These values describe token efficiency at selected validation thresholds, not wall-clock speedup or energy efficiency. Token counts are taken from the first observed evaluation point at or below each target, without interpolation.

4.4 Training-phase behavior

The initial training phase covers approximately 0–1.5B tokens. The hybrid's reported BPB is 0.028 lower than the baseline at the matched 1.5B-token point.

During extended training, an optimizer-state reset at approximately 1.5B tokens produced a visible BPB spike. The reported curves show the hybrid retaining its lead through the extended phase.

The reset is part of the experimental history and should be considered when interpreting the training curves.

5. GDN-2 kernel benchmark

The custom atomic_gdn2 kernel is implemented in JAX and Pallas for TPU execution.

Configuration	Value
Target	TPU v5e
Diagonal mode	BTL
Compute mode	bf16x3
bt	128
bc	64
bs2	32
mb	8
B3 dot mode	bf16x3
B4 dot mode	highest
Log-decay domain	[-5, 0]

5.1 Reported benchmark results

Heads	Sequence length	Speedup
2	2048	193×
2	4096	207×
4	2048	300×
4	4096	322×
6	2048	468×
6	4096	448×

The benchmark measures forward and backward execution against a JAX associative_scan implementation with batch size 4 and head dimension 128 on a single TPU v5e chip.

The maximum reported speedup is 468× at six heads and a sequence length of 2048. This is a kernel microbenchmark result, not an end-to-end model-training speedup.

Benchmark reproduction is documented in notebooks/speedup_scaling.ipynb.

6. Reproducibility and checkpoints

6.1 Repository contents

The Hugging Face repository Akseleu/Atomic-light contains two published checkpoints with per-checkpoint model cards and configuration files:

Akseleu/Atomic-light/
├── README.md                    root model card
├── .gitattributes               LFS rules for *.safetensors
├── baseline/
│   ├── README.md                334M, 28 × GDN-2
│   ├── config.json              architecture
│   ├── training_args.json       hyperparameters + metrics
│   └── model.safetensors        334M params, ~1.34 GB
└── hybrid/
    ├── README.md                342M, GDN-2 + MLA + DAR
    ├── config.json
    ├── training_args.json
    └── model.safetensors        342M params, ~1.37 GB

Approximate checkpoint sizes are based on parameter counts and fp32 storage. Checkpointing during training used Orbax with a configured interval of 5000 steps and a maximum of two retained checkpoints; this describes the training configuration, not the artifacts retained in the public repository.

The training implementation (train_hybrid.py) and the kernel benchmark notebook (notebooks/speedup_scaling.ipynb) are hosted at github.com/Akseleu-J/atomic-ops.

6.2 Loading the hybrid checkpoint

from huggingface_hub import snapshot_download
from safetensors.flax import load_file

repo = snapshot_download(
    "Akseleu/Atomic-light",
    allow_patterns=["hybrid/*"],
)

params = load_file(
    f"{repo}/hybrid/model.safetensors"
)

The safetensors file contains flattened parameter keys, including names such as Embed_0_embedding and block_0_layer_0_..._kernel.

Loading the file returns a flattened parameter dictionary. Reconstructing the Flax parameter tree and running inference requires the corresponding model definition and execution code in train_hybrid.py.

6.3 Environment

JAX          0.11.2
jaxlib       0.11.2
libtpu       0.0.48
Flax         0.12.9
Optax        0.2.8
Hardware     TPU v5e-8
Hub          Hugging Face Hub
Checkpoint   safetensors, Orbax

7. Limitations

1. Single seed. Each configuration was trained with one reported seed. Three-seed evaluation is planned to estimate run-to-run variance.
2. Incomplete baseline training. The baseline stopped at approximately 3.71B tokens rather than completing the planned 4.497B-token budget.
3. Different batch sizes. The baseline uses batch size 48; the hybrid uses batch size 32.
4. Different optimizer-step counts. The models do not have matched numbers of parameter updates at the final checkpoints.
5. Different compute budgets. Token matching does not imply FLOP matching.
6. Optimizer-state reset. The extended phase includes a reset at approximately 1.5B tokens.
7. No component ablations. The individual contributions of MLA and DAR have not been isolated.
8. No external model baselines. Mamba-2, RWKV-7, and GLA have not been evaluated under the same experimental setup.
9. Kernel and model performance differ. Microbenchmark speedups do not establish equivalent end-to-end training speedups.
10. Evaluation details. The validation split, checkpoint-selection rule, and threshold-crossing procedure should be specified with the evaluation code.
11. Research artifacts. Deployment behavior, generation quality, and production readiness have not been comprehensively evaluated.

8. Conclusion

Atomic-GDN2 evaluates a hybrid sequence-modeling architecture that combines GDN-2 recurrence, MLA, and hierarchical residual mixing.

At the reported matched-token checkpoint of 1.5B tokens, the hybrid reaches 1.0671 BPB compared with 1.0949 BPB for the baseline, a difference of approximately 0.028 BPB.

The hybrid also reaches the reported intermediate BPB thresholds with fewer training tokens. Its final checkpoint was obtained after approximately 4.50B tokens, whereas the baseline was interrupted at approximately 3.71B tokens. The final checkpoints therefore do not represent equal-budget training runs.

The results motivate further evaluation using multiple seeds, component-level ablations, matched compute budgets, and a completed baseline run.

9. Artifacts

• Model repository: https://huggingface.co/Akseleu/Atomic-light
• Kernel implementation: https://github.com/Akseleu-J/atomic-ops
• Kernel benchmark: notebooks/speedup_scaling.ipynb
• Training implementation: train_hybrid.py
• arXiv identifier: to be added after submission.

10. Citation

@misc{akseleu2026atomicgdn2,
  title = {Atomic-GDN2: Gated DeltaNet-2 with Multi-head Latent Attention and Hierarchical Block-Delta Residuals},
  author = {Akseleu Omirbay},
  year = {2026},
  howpublished = {\url{https://huggingface.co/Akseleu/Atomic-light}},
  note = {Preprint identifier to be added after arXiv submission}
}