AMOR-Mamba2-1.5B

Paper · Code · Files · All models

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.

For a smaller download, see the 180M and 440M models below.

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

Detailed settings are in training_config.json; training_summary.json records the completed run.

Results

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

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

Research comparison across all three scales; use the 1.5B 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.

Retrieval and long context

The 1.5B AMOR variants improve average retrieval over their recurrent backbones; serial hybrids retain stronger six-task retrieval averages. LongBench results remain close to the recurrent backbones, with lower scores for the serial hybrids in this experimental setting.

Left: six-task retrieval mean. Right: averages of three single-needle variants at each context length; lengths are not pooled.

Unweighted mean of 14 task-native scores. Whiskers are benchmark standard errors propagated under task independence, not training-seed variability.

<details> <summary>Per-task retrieval and LongBench tables</summary>

Retrieval outputs do not store standard errors.

Scores are on a 0–100 scale. The mean SE is propagated under task independence.

</details>

Per-task values are also available in 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). Then, in a fresh Python environment:

python -m pip install huggingface_hub
hf download FlyinGodzilla/AMOR-Mamba2-1.5B --local-dir AMOR-Mamba2-1.5B
cd AMOR-Mamba2-1.5B
python -m pip install "transformers==5.0.0" safetensors einops
python -m pip install "mamba-ssm[causal-conv1d]" --no-build-isolation
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. Mamba2 kernels: mamba-ssm 2.3.0 and causal-conv1d 1.6.0. 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 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 8 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 records provenance and actual validation; manifest.json lists file hashes.

Model family

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.

Citation

@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}
}