README.md

---
language: en
library_name: pytorch
pipeline_tag: text-generation
license: mit
tags:
- amor
- custom-code
- pretrained
- research
- arxiv:2602.13215
- gateddeltanet
datasets:
- HuggingFaceFW/fineweb-edu
---

![AMOR-Gated DeltaNet 440M — When to think fast and slow?](assets/title.svg)

[Paper](https://arxiv.org/abs/2602.13215) · [Code](https://github.com/HaoranZhengRaul/AMOR) · [Files](https://huggingface.co/FlyinGodzilla/AMOR-GatedDeltaNet-440M/tree/main) · [All models](#model-family)

![AMOR architecture: recurrent backbone and entropy-gated refinement](assets/architecture.png)

*Block 0 receives the backbone's final normalized output. Blocks 1 and 2 apply
their own pre-normalization. A shared, embedding-tied LM head supplies all three
entropy evaluations and the final output prediction.*

**A complete recurrent backbone, refined by three entropy-gated attention blocks.**

Looking for the larger model? [AMOR-Gated DeltaNet 1.5B](https://huggingface.co/FlyinGodzilla/AMOR-GatedDeltaNet-1.5B).

## The idea

AMOR (**Adaptive Metacognitive Output Router**) uses the model's predictive
uncertainty to decide where attention should refine its recurrent representation.
Each appended block reads normalized output entropy, compares it with a frozen
threshold at inference, and adds an attention update only where its gate fires.

The recurrent backbone and attention blocks are pretrained jointly from scratch.
“Post-hoc” describes where the attention blocks sit, after the complete backbone.
Keys and values come from residual-stream projections and are retained at every
position for future retrieval.


This model uses the **native entropy gate**. Optional distilled
deployment routers are included and enabled with --router. Training uses dense
attention with an output mask, so gating does not reduce training attention FLOPs.

### Fitted inference router
The package includes three fitted width-512 SiLU predictors with linear skips.
They estimate entropy from the same normalized residual and use the original
frozen thresholds. Router weights and computation stay fp32 under autocast.
Calibration is included in the fitted weights. Native entropy gating remains
the default; --router selects the fitted predictors for full-sequence inference,
prefill and cached decoding. Base-model training uses the native entropy gate.

In Python, use load_model(directory, device="cuda", use_router=True).
The loader verifies the base-weight checksum and matching router configuration.
The figures and benchmark tables below describe native entropy-gate models.

## This model

| Property | Value |
|---|---|
| Backbone | Gated DeltaNet |
| Parameters | 442,133,056 |
| Recurrent mixer–MLP layers | 24 |
| Model width | 1024 |
| Appended attention blocks | 3 |
| Training corpus | FineWeb-Edu |
| Training tokens | 8,973,017,088 |
| Training context | 3,072 tokens |
| Final training step | 18,255 |
| Seed | 42 |
| Tokenizer vocabulary | 128,256 |
| Model type | Pretrained base model; not instruction-tuned |
| Weight format | Safetensors; original weight precision |

Detailed settings are in [training_config.json](training_config.json);
[training_summary.json](training_summary.json) records the completed run.

## Results

This model's canonical eight-task common-sense mean is **40.5%**.
The score uses raw accuracy on LAMBADA, HellaSwag, PIQA, ARC-Easy, ARC-Challenge,
WinoGrande, OpenBookQA and TruthfulQA-mc2.

| Task | Accuracy (%) |
|---|---:|
| LAMBADA | 30.2 |
| HellaSwag | 31.8 |
| PIQA | 65.0 |
| ARC-Easy | 56.0 |
| ARC-Challenge | 24.6 |
| WinoGrande | 51.3 |
| OpenBookQA | 23.8 |
| TruthfulQA-mc2 | 40.9 |
| **Eight-task mean** | **40.5** |

Per-task entries are point estimates; the figure reports evaluation-repeat spread.

![Common-sense reasoning across three model scales](assets/common_sense.png)

*Research comparison across all three scales; use the **440M** panel for this
model. Whiskers show the sample SD of five reevaluations of each fixed
model, centered on its canonical score. They measure evaluation
repeatability, not independent-training variation. Axes are zoomed per panel.*



Per-task values are also available in [benchmark_results.json](benchmark_results.json).

## Download and run

These are custom PyTorch models. Use the supplied loader and generation script;
this repository is not a Transformers `AutoModel.from_pretrained` integration,
a hosted chat endpoint, or a GGUF/PocketPal package.

A Linux machine with a compatible NVIDIA GPU is the intended accelerated setup.
Install a CUDA-enabled PyTorch build appropriate for your driver first
([PyTorch installation](https://pytorch.org/get-started/locally/)).
Then, in a fresh Python environment:

```bash
python -m pip install huggingface_hub
hf download FlyinGodzilla/AMOR-GatedDeltaNet-440M --local-dir AMOR-GatedDeltaNet-440M
cd AMOR-GatedDeltaNet-440M
python -m pip install "transformers==5.0.0" safetensors einops
python -m pip install flash-linear-attention
python generate.py --prompt "The Eiffel Tower is" --max-new-tokens 32
python generate.py --router --prompt "The Eiffel Tower is" --max-new-tokens 32
```

For the tested CPU reference path, use PyTorch 2.6.0 and Transformers 4.51.3.
Mamba2 has a slow CPU fallback (`--device cpu`); Gated DeltaNet requires its
compatible accelerated kernels. GPU kernel versions must match the PyTorch/CUDA
environment. CUDA smoke tests used PyTorch 2.6.0+cu124, Transformers 5.0.0,
Triton 3.2.0 and safetensors 0.7.0.
Gated DeltaNet kernels: the imported fla module reported version 0.4.2.
These are actual imported versions from the tested environment; the installation
commands above have not been checked in a fresh environment.

`generate.py` performs cached greedy text completion. It is intentionally a
plain-prompt example: these models were not trained as chat assistants.
The tokenizer is referenced from
[teknium/Llama-3.1-AlternateTokenizer](https://huggingface.co/teknium/Llama-3.1-AlternateTokenizer)
and pinned in `config.json` to the release-packaging revision. The original
training-time tokenizer revision was not recorded.

The base weights are split into **2 standard safetensors files** for easier
downloading. Keep every model shard and `model.safetensors.index.json` in the
same folder. The download command retrieves them together; the loader reads
them automatically. Tensor values and precision are unchanged. The optional
routers remain separate in `router.safetensors` and `router_config.json`.

### Verification
Both entropy-gate and fitted-router modes passed full-forward, prefill, cached
decoding and the generation CLI on H100 NVL. Tested contexts were 16, 64 and
128 tokens (batches 1 and 2 for prefill); cached decoding checked four steps
at batch 1. The maximum cached/full-forward absolute logit difference was
0.125 (rtol 0.02, atol 0.1). Default-mode forward logits matched the prior
release exactly. Router predictions matched the research implementation on
the tested inputs, and router mode used the LM head only for final predictions.
Seventeen tests covered router precision and attachment, sharded-weight
integrity, loader precision and safe downloads.

Every exported tensor was compared exactly with the original trained weights,
including the entropy gate's frozen buffers. The tied embedding/head matrix is
stored once and restored by the loader. Optimizer state is excluded.
[Export metadata](export_metadata.json) records provenance and actual validation;
[manifest.json](manifest.json) lists file hashes.

## Model family
| Backbone | 180M | 440M | 1.5B |
|---|---|---|---|
| Mamba2 | [180M](https://huggingface.co/FlyinGodzilla/AMOR-Mamba2-180M) | [440M](https://huggingface.co/FlyinGodzilla/AMOR-Mamba2-440M) | [1.5B](https://huggingface.co/FlyinGodzilla/AMOR-Mamba2-1.5B) |
| Gated DeltaNet | [180M](https://huggingface.co/FlyinGodzilla/AMOR-GatedDeltaNet-180M) | [440M](https://huggingface.co/FlyinGodzilla/AMOR-GatedDeltaNet-440M) | [1.5B](https://huggingface.co/FlyinGodzilla/AMOR-GatedDeltaNet-1.5B) |

All six models use the three-block canonical AMOR layout. Choose a backbone
and scale to match your experiment; larger models require more memory.
The cards use a snapshot of the research figures and do not automatically track
later manuscript revisions.

## Intended use and limitations
Research on recurrent–attention hybrids, uncertainty-based routing, and pretrained
language-model behavior. Outputs can be incorrect and may inherit training-data
biases. Entropy measures predictive uncertainty; it cannot detect every confidently
wrong prediction. Inference fire rates depend on the workload, and attention KV
storage still grows with context. Training context is 3,072 tokens; long-context
experiments do not establish reliable performance on every domain.

## License
The code and model weights are released under the MIT license; see [LICENSE](https://huggingface.co/FlyinGodzilla/AMOR-GatedDeltaNet-440M/blob/main/LICENSE).

## Citation
```bibtex
@article{zheng2026amor,
  title={When to Think Fast and Slow? AMOR: Adaptive Entropy Gate for Hybrid Models},
  author={Zheng, Haoran and Shani, Chen},
  journal={arXiv preprint arXiv:2602.13215},
  year={2026}
}
```