OpenMythos Documentation

A theoretical reconstruction of the Claude Mythos architecture, built from first principles using the available research literature.

RAW Doc

File: README.md

OpenMythos

Disclaimer: OpenMythos is an independent, community-driven theoretical reconstruction based solely on publicly available research and speculation. It is not affiliated with, endorsed by, or connected to Anthropic or any of their proprietary systems.

OpenMythos is an open-source, theoretical implementation of the Claude Mythos model. It implements a Recurrent-Depth Transformer (RDT) with three stages: Prelude (transformer blocks), a looped Recurrent Block (up to max_loop_iters), and a final Coda. Attention is switchable between MLA and GQA, and the feed-forward uses a sparse MoE with routed and shared experts ideal for exploring compute-adaptive, depth-variable reasoning.

Installation

bash
pip install open-mythos

#uv pip install open-mythos

To enable Flash Attention 2 in GQAttention (requires CUDA and build tools):

bash
pip install open-mythos[flash]

Usage

python
import torch
from open_mythos.main import OpenMythos, MythosConfig


attn_type = "mla"  # or "gqa"

base = {
    "vocab_size": 1000,
    "dim": 256,
    "n_heads": 8,
    "max_seq_len": 128,
    "max_loop_iters": 4,
    "prelude_layers": 1,
    "coda_layers": 1,
    "n_experts": 8,
    "n_shared_experts": 1,
    "n_experts_per_tok": 2,
    "expert_dim": 64,
    "lora_rank": 8,
    "attn_type": attn_type,
}

if attn_type == "gqa":
    cfg = MythosConfig(**base, n_kv_heads=2)
else:
    cfg = MythosConfig(
        **base,
        n_kv_heads=8,
        kv_lora_rank=32,
        q_lora_rank=64,
        qk_rope_head_dim=16,
        qk_nope_head_dim=16,
        v_head_dim=16,
    )

model = OpenMythos(cfg)
total = sum(p.numel() for p in model.parameters())
print(f"\n[{attn_type.upper()}] Parameters: {total:,}")

ids = torch.randint(0, cfg.vocab_size, (2, 16))
logits = model(ids, n_loops=4)
print(f"[{attn_type.upper()}] Logits shape: {logits.shape}")

out = model.generate(ids, max_new_tokens=8, n_loops=8)
print(f"[{attn_type.upper()}] Generated shape: {out.shape}")

A = model.recurrent.injection.get_A()
rho = torch.linalg.eigvals(A).abs().max().item()
print(
    f"[{attn_type.upper()}] Spectral radius ฯ(A) = {rho:.4f} (must be < 1)"
)

Model Variants

Pre-configured scales from 1B to 1T parameters:

python
from open_mythos import (
    mythos_1b,
    mythos_3b,
    mythos_10b,
    mythos_50b,
    mythos_100b,
    mythos_500b,
    mythos_1t,
    OpenMythos,
)

cfg = mythos_7b()  # returns a MythosConfig
model = OpenMythos(cfg)

total = sum(p.numel() for p in model.parameters())
print(f"Parameters: {total:,}")
Variant dim Experts expert_dim Loop iters Context Max output
mythos_1b 2048 64 2048 16 4k 4k
mythos_3b 3072 64 4096 16 4k 4k
mythos_10b 4096 128 5632 24 8k 4k
mythos_50b 6144 256 9728 32 8k 4k
mythos_100b 8192 256 13568 32 1M 128k
mythos_500b 12288 512 23040 48 1M 128k
mythos_1t 16384 512 34560 64 1M 128k

Training

The training script for the 3B model on FineWeb-Edu is at training/3b_fine_web_edu.py.

Single GPU:

bash
python training/3b_fine_web_edu.py

Multi-GPU (auto-detects GPU count):

bash
torchrun --nproc_per_node=$(python -c "import torch; print(torch.cuda.device_count())") training/3b_fine_web_edu.py

Key design choices:

Feature Detail
Optimizer AdamW
Dataset HuggingFaceFW/fineweb-edu (sample-10BT by default, swap to sample-100BT or default for full run)
Tokenizer openai/gpt-oss-20b via MythosTokenizer
Parallelism PyTorch DDP via torchrun, sharded streaming dataset
Precision bfloat16 on H100/A100, float16 + GradScaler on older GPUs
Schedule Linear warmup (2000 steps) โ†’ cosine decay
Target 30B tokens (~Chinchilla-adjusted for looped architecture)

Documentation

Page Description
docs/open_mythos.md Full API reference for the OpenMythos class โ€” constructor, forward, generate, all sub-modules, configuration reference, and usage examples
docs/datasets.md Recommended training datasets with token budget guidance per model size

The Central Hypothesis

Claude Mythos is suspected to be a Recurrent-Depth Transformer (RDT) โ€” also called a Looped Transformer (LT). Rather than stacking hundreds of unique layers, a subset of layers is recycled and run through multiple times per forward pass. Same weights. More loops. Deeper thinking.

This is not chain-of-thought. There is no intermediate token output. All of this reasoning happens silently, inside a single forward pass, in continuous latent space.


Architecture

A looped transformer divides its layers into three functional blocks:

text
Input
  โ†“
[Prelude P]        โ€” standard transformer layers, run once
  โ†“
[Recurrent Block R] โ€” looped T times
  โ†‘_______โ†“         (hidden state h updated each loop with input injection e)
  โ†“
[Coda C]           โ€” standard transformer layers, run once
  โ†“
Output

The recurrent block update rule at each loop step t:

text
h_{t+1} = Aยทh_t + Bยทe + Transformer(h_t, e)

Where:

  • h_t is the hidden state after loop t
  • e is the encoded input (from the Prelude), injected at every loop
  • A and B are learned injection parameters
  • The Transformer blocks apply attention and MLP as usual

The injection of e at every step is what prevents the model from drifting โ€” it keeps the original input signal alive throughout the entire recurrence depth.

The full implementation is in open_mythos/main.py. See the OpenMythos class reference for a detailed API walkthrough, configuration options, and usage examples.

Attention Implementations

The attention layer is switchable via cfg.attn_type:

Option Class Description
"gqa" GQAttention Grouped Query Attention (Ainslie et al., 2023) โ€” fewer KV heads than Q heads (n_kv_heads < n_heads), reducing KV-cache memory by n_heads / n_kv_heads. Uses Flash Attention 2 (Dao et al., 2023) when flash-attn>=2.8.3 is installed: GQA is handled natively (no KV head expansion), I/O-bound-optimal, with a transparent fallback to manual scaled dot-product attention when the package is absent.
"mla" MLAttention Multi-Latent Attention (DeepSeek-V2) โ€” caches a compressed KV latent (kv_lora_rank) rather than full K/V, with split RoPE / no-RoPE head dims for position-aware compression.

RoPE is applied to Q and K before caching, so cached values do not need to be re-rotated on retrieval.


Why This Explains Mythos

1. Systematic Generalization

Vanilla transformers fail to combine knowledge in ways they have never seen during training. Looped transformers pass this test. The ability emerges through a three-stage grokking process:

  1. Memorization โ€” model fits training distribution
  2. In-distribution generalization โ€” model handles known compositions
  3. Systematic generalization โ€” model handles novel compositions OOD, abruptly and suddenly

This is why Mythos feels qualitatively different from other models on novel questions โ€” the capability phase-transitions in, rather than emerging gradually.

2. Depth Extrapolation

Train on 5-hop reasoning chains. Test on 10-hop. Vanilla transformer fails. Looped transformer succeeds โ€” by running more inference-time loops. This maps directly to the observation that Mythos handles deeply compositional problems (multi-step math, long-horizon planning, layered arguments) without explicit chain-of-thought.

More loops at inference = deeper reasoning chains = harder problems solved.

3. Latent Thoughts as Implicit Chain-of-Thought

Each loop iteration is the functional equivalent of one step of chain-of-thought, but operating in continuous latent space rather than token space. A looped model running T loops implicitly simulates T steps of CoT reasoning. This has been formally proven (Saunshi et al., 2025).

Furthermore, continuous latent thoughts โ€” unlike discrete token outputs โ€” can encode multiple alternative next steps simultaneously. This allows something closer to breadth-first search over the reasoning space, rather than a single committed reasoning path. The model is effectively exploring many possible directions inside each forward pass before converging.

4. No Parameter Explosion

A looped model with k layers run L times achieves the quality of a kL-layer non-looped model, with only k layers worth of parameters. For Mythos-scale deployments, this matters enormously:

  • Memory footprint does not grow with reasoning depth
  • Inference-time compute scales with loop count, not model size
  • This makes deeper reasoning "free" in terms of parameters

The Stability Problem (and How It Was Likely Solved)

Training looped models is notoriously unstable. Two failure modes dominate:

  • Residual explosion โ€” the hidden state h_t grows unboundedly across loops
  • Loss spikes โ€” training diverges suddenly due to large spectral norms in injection parameters

The Dynamical Systems View

Recast looping as a discrete linear time-invariant (LTI) dynamical system over the residual stream. Ignoring the nonlinear Transformer contribution, the recurrence becomes:

text
h_{t+1} = Aยทh_t + Bยทe

For this LTI system, stability is governed entirely by the spectral radius of A:

  • ฯ(A) < 1 โ†’ stable, convergent
  • ฯ(A) โ‰ฅ 1 โ†’ unstable, divergent

Empirically, every divergent training run learns ฯ(A) โ‰ฅ 1. Every convergent run maintains ฯ(A) < 1.

The Fix

Constrain the injection parameters so that stability is guaranteed by construction:

  1. Parameterize A as a continuous negative diagonal matrix
  2. Discretize using ZOH/Euler schemes: A_discrete = exp(ฮ”t ยท A_continuous)
  3. Enforce negativity via A := Diag(-exp(log_A)) with a learned scalar ฮ”t
  4. This ensures ฯ(A) < 1 always holds, regardless of learning rate or batch noise

The result: the looped model becomes significantly more robust to hyperparameter selection and trains cleanly even at high learning rates. This is the Parcae architecture (Prairie et al., 2026), and it represents the most likely class of solution Anthropic used to make Mythos trainable.


Scaling Laws for Looped Models

Parcae establishes the first predictable scaling laws for looped training:

  • Training: For a fixed FLOP budget with fixed parameters, increasing mean recurrence and reducing token count yields a lower loss than training with minimal loops on more data. Optimal recurrence and optimal token count both follow power laws with consistent exponents across scales.
  • Inference: More test-time loops improves quality following a predictable, saturating exponential decay โ€” gains are real but diminishing. This mirrors the inference-time scaling of chain-of-thought.

At 770M parameters, a looped model achieves the downstream quality of a 1.3B fixed-depth Transformer trained on the same data โ€” roughly half the parameters for the same quality.

Applied to Mythos: if trained under these scaling laws, Mythos could be dramatically more parameter-efficient than it appears, with a large fraction of its apparent "capability" coming from loop depth rather than raw parameter count.


The Loop Index Embedding Hypothesis

A key open question is whether the looped block behaves identically on every iteration, or whether it can learn to do different things at different loop depths.

Without any positional signal across loops, the same weights must handle both early-stage pattern matching and late-stage refinement โ€” a tight constraint. A RoPE-like embedding of the loop index injected alongside the input at each step would allow the same parameters to implement functionally distinct operations across iterations, much like how RoPE allows the same attention heads to behave differently at different sequence positions.

If Mythos uses this technique, each loop is not a repetition โ€” it is a distinct computational phase, all sharing weights but operating in different representational regimes. This would substantially increase the expressiveness of the recurrent block without increasing parameter count.


The Overthinking Problem

More loops is not always better. Beyond a certain depth, excessive recurrence degrades predictions โ€” the hidden state drifts past the solution and into noise. This is the "overthinking" failure mode.

The original Universal Transformer (Dehghani et al., 2018) addressed this with an Adaptive Computation Time (ACT) halting mechanism: a learned scalar per position that dynamically decides when to stop looping. Positions that are harder to process receive more computation; simple tokens halt early.

Mythos almost certainly has some version of this. The model cannot naively run the maximum number of loops on every input โ€” it needs a learned signal for when the answer has converged. The ACT mechanism also makes the model Turing-complete under certain assumptions, which has theoretical implications for the class of problems it can solve.


Mixture of Experts โ€” Suspected for Large Parameter Counts

The looped transformer explains the depth of Mythos's reasoning, but not the breadth. Handling wildly different domains โ€” code, math, literature, science, law โ€” with the same weights requires Mixture of Experts (MoE). The suspected design replaces every FFN in the Recurrent Block with a fine-grained MoE layer: each FFN is split into many small experts (1/m the normal size), a router selects the top-mK of them per token via learned affinity scores, and a small number of shared experts are always activated regardless of routing to absorb common cross-domain knowledge โ€” syntax, basic reasoning, general context โ€” that would otherwise be redundantly learned by every routed expert. Routing collapse is prevented through a bias term on the router logits adjusted dynamically during training, keeping load balanced across experts without distorting the loss signal.

As the hidden state h_t evolves across loop iterations, the router may select different expert subsets at each depth, making every loop computationally distinct despite shared weights. MoE provides breadth; looping provides depth. If the activation ratio is ~5%, Mythos could hold hundreds of billions of total parameters while activating only a small fraction per token โ€” the true parameter count, if ever disclosed, would be a storage number, not a compute number.


The Memorization-Reasoning Tradeoff

Looped models exhibit an interesting dichotomy: looping improves reasoning but can hurt memorization. The recurrent structure is optimized for iterative composition โ€” running a reasoning chain forward โ€” but does not inherently improve the storage of rote facts.

This maps to an observable characteristic of Mythos: it reasons exceptionally well about novel problems it has never seen, but its factual recall can be inconsistent. The architecture is structurally biased toward composition over memorization.

Looping-based regularization (Saunshi et al., 2025) can be used to balance this tradeoff during training โ€” applying stronger looping constraints for reasoning tasks while relaxing them for retrieval tasks.


Parameter Reuse via LoRA Adaptation

A complementary approach from Relaxed Recursive Transformers (Bae et al., 2024): rather than requiring fully identical weights at every loop, add a small depth-wise LoRA module at each iteration. This preserves the compactness of weight sharing while allowing each loop to adapt its behavior slightly.

The result:

  • Each loop shares a large common weight matrix (the recursive base)
  • A small rank-r adaptation matrix shifts behavior per iteration depth
  • The total parameter overhead is minimal

This bridges the gap between pure weight-tying (maximally parameter-efficient, less expressive) and fully distinct layers (maximally expressive, no parameter savings). Mythos likely sits somewhere on this spectrum.


Continuous Depth-wise Batching

A downstream consequence of the recursive architecture: Continuous Depth-wise Batching. Because all tokens share the same recurrent block, the model can exit the loop at different depths for different tokens or sequences โ€” processing easy inputs quickly and hard inputs with more iterations, all within the same batch.

Theoretical analysis suggests 2-3x improvements in inference throughput. For a deployed model like Mythos serving many users simultaneously, this would be a substantial efficiency gain.


Summary: What Mythos Probably Is

Property Description
Architecture Recurrent-Depth Transformer (Prelude + Looped Recurrent Block + Coda)
FFN layer Suspected MoE โ€” fine-grained experts + always-on shared experts
Parameter count Very large total; small fraction activated per token (~5% estimate)
Reasoning mechanism Implicit multi-hop via iterative latent updates โ€” no token output between steps
Inference-time scaling More loops = deeper reasoning, following predictable exponential decay
Training stability LTI-constrained injection parameters with spectral radius < 1
Loop differentiation Likely uses loop-index positional embedding (ร  la RoPE) per iteration
Halting Adaptive Computation Time or learned convergence criterion
Attention GQA (with optional Flash Attention 2) or MLA with compressed KV latent cache
Scaling law Optimal training scales looping and data together, not parameters alone
Reasoning vs. memory Structurally biased toward composition; memorization requires separate treatment
Deployment Continuous Depth-wise Batching enables variable compute per request

References

Twitter / X

Papers


Citation

If you use OpenMythos in your research or build on this work, please cite:

bibtex
@software{gomez2026openmythos,
  author    = {Kye Gomez},
  title     = {OpenMythos: A Theoretical Reconstruction of the Claude Mythos Architecture},
  year      = {2026},
  url       = {https://github.com/kyegomez/OpenMythos},
  note      = {Recurrent-Depth Transformer with MoE, MLA, LTI-stable injection, and ACT halting}
}

License

MIT License โ€” Copyright (c) 2026 Kye Gomez. See LICENSE for the full text.


File: requirements.txt

torch>=2.1.0
transformers>=4.40.0
datasets>=2.18.0
pytest>=7.0.0

optional โ€” enables Flash Attention 2 in GQAttention (requires CUDA + build tools)

flash-attn>=2.8.3


File: docs/datasets.md

Dataset HuggingFace Tokens License Use
FineWeb-Edu HuggingFaceFW/fineweb-edu 1.3T Apache 2.0 Primary pretraining
OpenHermes 2.5 teknium/OpenHermes-2.5 ~1M samples Apache 2.0 Instruction tuning (~5% mix)
OpenWebMath open-web-math/open-web-math 14.7B ODC-By Math/reasoning boost

Primary Pretraining

FineWeb-Edu

  • HuggingFace: HuggingFaceFW/fineweb-edu
  • Size: 1.3T tokens
  • License: Apache 2.0
  • Why: Web text filtered for educational quality. Outperforms The Pile, C4, and RefinedWeb on downstream benchmarks. Already deduplicated and cleaned.
  • Start with: sample-10BT to validate your pipeline, then sample-100BT or the full corpus for a serious run.

Supplementary

OpenHermes 2.5

  • HuggingFace: teknium/OpenHermes-2.5
  • Size: ~1M instruction samples
  • License: Apache 2.0
  • Why: High-quality instruction-following data. Mix in ~5% by token count on top of FineWeb-Edu to improve instruction following without degrading general capability.

OpenWebMath

  • HuggingFace: open-web-math/open-web-math
  • Size: ~14.7B tokens
  • License: ODC-By
  • Why: Math-focused web text. Add if you want stronger quantitative and symbolic reasoning. Particularly useful for the 10B+ variants where reasoning depth matters.

Token Budget Recommendations

Variant Chinchilla-optimal Recommended (looped)
1B ~20B tokens ~10โ€“15B tokens
3B ~60B tokens ~30โ€“40B tokens
10B ~200B tokens ~100โ€“150B tokens
50B+ ~1T+ tokens ~500B+ tokens

The looped architecture is more sample-efficient than a standard transformer โ€” same validation loss is reachable with fewer tokens due to faster convergence. The "recommended (looped)" column reflects this and is based on the Tiny Shakespeare result where OpenMythos reached equivalent loss ~2.5ร— faster than nanoGPT.


File: docs/open_mythos.md

`OpenMythos` โ€” Class Reference

Module: open_mythos.main
Base class: torch.nn.Module


Overview

OpenMythos is the top-level model class implementing the Recurrent-Depth Transformer (RDT) architecture described in the OpenMythos hypothesis. It assembles three functional stages โ€” Prelude, Recurrent Block, and Coda โ€” into a complete autoregressive language model.

text
Input token IDs  (B, T)
        โ†“
   [Embedding]          token index โ†’ dim-dimensional vector
        โ†“
   [Prelude]            prelude_layers ร— standard TransformerBlock  (run once)
        โ†“
   [Recurrent Block]    one TransformerBlock looped T times
        โ†‘___________โ†“   h_{t+1} = Aยทh_t + Bยทe + Transformer(h_t, e)
        โ†“
   [Coda]               coda_layers ร— standard TransformerBlock  (run once)
        โ†“
   [RMSNorm โ†’ LM head]
        โ†“
Output logits  (B, T, vocab_size)

Every architectural choice in OpenMythos can be configured through a single MythosConfig dataclass passed at construction.


`MythosConfig`

python
@dataclass
class MythosConfig

All hyperparameters for the model are stored in this single frozen-style dataclass. Pass an instance to OpenMythos.__init__.

Core fields

Field Type Default Description
vocab_size int 32000 Token vocabulary size; sets the embedding and LM head dimension
dim int 2048 Model hidden dimension โ€” the width of the residual stream throughout
n_heads int 16 Number of query attention heads
n_kv_heads int 4 Number of key/value heads (GQA only); n_heads // n_kv_heads Q heads share each KV pair
max_seq_len int 4096 Maximum sequence length; RoPE frequencies are precomputed up to this length
max_loop_iters int 16 Default recurrent loop depth T at inference. Can be overridden per call
prelude_layers int 2 Number of standard transformer blocks run once before the recurrent loop
coda_layers int 2 Number of standard transformer blocks run once after the recurrent loop

Attention fields

attn_type selects between two complete attention implementations. All other attention fields are implementation-specific.

Field Type Default Description
attn_type str "mla" "gqa" for Grouped Query Attention; "mla" for Multi-Latent Attention
kv_lora_rank int 512 [MLA only] Compressed KV latent rank stored in the cache instead of full K and V
q_lora_rank int 1536 [MLA only] Compressed Q latent rank
qk_rope_head_dim int 64 [MLA only] Per-head dimension receiving RoPE positional encoding
qk_nope_head_dim int 128 [MLA only] Per-head dimension without positional encoding
v_head_dim int 128 [MLA only] Per-head value dimension

GQA vs MLA: GQA reduces KV cache by having fewer KV heads than Q heads (factor of n_heads / n_kv_heads). MLA achieves a much larger reduction by caching a low-rank KV latent (kv_lora_rank) and the RoPE keys (n_heads ร— qk_rope_head_dim), then reconstructing full K and V on the fly. At production scale MLA yields roughly 10โ€“20ร— smaller KV cache than standard attention.

MoE FFN fields

The Mixture-of-Experts FFN is used exclusively inside the Recurrent Block. Prelude and Coda use a dense SwiGLU FFN.

Field Type Default Description
n_experts int 64 Total number of routed expert FFNs
n_shared_experts int 2 Always-active shared experts; absorb common cross-domain patterns
n_experts_per_tok int 4 Top-K routed experts selected per token by the router
expert_dim int 512 Hidden dimension inside each fine-grained routed expert

Approximately n_experts_per_tok / n_experts = 6.25% of routed expert parameters are activated per token, plus all shared expert parameters.

Stability and adaptation fields

Field Type Default Description
act_threshold float 0.99 ACT cumulative halting threshold; loop exits per-position once this is exceeded
rope_theta float 500000.0 RoPE base frequency (LLaMA-3 default; higher = slower frequency decay over sequence positions)
lora_rank int 16 Rank of the depth-wise LoRA adapter applied inside each loop iteration

Constructor

python
OpenMythos(cfg: MythosConfig)

Builds all sub-modules, precomputes RoPE frequency buffers, and runs weight initialization.

What happens internally:

  1. nn.Embedding(vocab_size, dim) โ€” token embedding table, weight-tied with the LM head.
  2. RoPE buffers โ€” freqs_cis (for GQA, dim = dim // n_heads) and freqs_cis_mla (for MLA, dim = qk_rope_head_dim) are precomputed once and registered as non-parameter buffers. The correct buffer is selected at forward time based on cfg.attn_type.
  3. prelude โ€” nn.ModuleList of prelude_layers TransformerBlock instances with dense SwiGLU FFN.
  4. recurrent โ€” a single RecurrentBlock containing one TransformerBlock (with MoE FFN), LTIInjection, ACTHalting, and LoRAAdapter.
  5. coda โ€” nn.ModuleList of coda_layers TransformerBlock instances with dense SwiGLU FFN.
  6. RMSNorm(dim) applied before the LM head.
  7. nn.Linear(dim, vocab_size, bias=False) LM head with weights tied to the embedding.
  8. All nn.Linear and nn.Embedding weights initialized from N(0, 0.02).

Example:

python
from open_mythos.main import OpenMythos, MythosConfig

cfg = MythosConfig(
    vocab_size=32000,
    dim=2048,
    n_heads=16,
    n_kv_heads=4,
    max_loop_iters=16,
    attn_type="mla",
)
model = OpenMythos(cfg)
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

`forward`

python
def forward(
    self,
    input_ids: torch.Tensor,
    n_loops: Optional[int] = None,
    kv_cache: Optional[dict] = None,
) -> torch.Tensor

Single forward pass through the full Prelude โ†’ Recurrent Block โ†’ Coda pipeline.

Parameters

Parameter Type Description
input_ids Tensor (B, T) Batch of token index sequences. B = batch size, T = sequence length
n_loops int | None Recurrent loop depth for this call. Defaults to cfg.max_loop_iters. Pass a higher value at inference to extrapolate to harder problems (depth extrapolation property).
kv_cache dict | None If provided, keys and values are accumulated here for autoregressive decoding. Pass {} on the first decode step and reuse the same dict across steps. Pass None for training or full-context inference.

Returns

Tensor (B, T, vocab_size) โ€” raw (unnormalized) logits over the vocabulary for each position.

Behavior walkthrough

text
1. Embed:     x = embedding(input_ids)              # (B, T, dim)
2. Select RoPE buffer:
     if attn_type == "mla": use freqs_cis_mla[:T]
     else:                   use freqs_cis[:T]
3. Build causal mask (upper-triangular -inf):
     if T > 1: mask = _causal_mask(T, device)
     else:     mask = None  (single-token decode step)
4. Prelude:
     for i, layer in prelude:
         x = layer(x, freqs_cis, mask, kv_cache, f"prelude_{i}")
5. Freeze encoded input:
     e = x                                          # (B, T, dim)
6. Recurrent loop:
     x = recurrent(x, e, freqs_cis, mask, n_loops, kv_cache)
7. Coda:
     for i, layer in coda:
         x = layer(x, freqs_cis, mask, kv_cache, f"coda_{i}")
8. Project:   logits = lm_head(norm(x))             # (B, T, vocab_size)

Step 5 (freeze e) is the key architectural invariant: the encoded input e is captured after the Prelude and injected at every loop iteration unchanged. This prevents the hidden state from drifting away from the original input signal regardless of loop depth.

Training example

python
import torch
from open_mythos.main import OpenMythos, MythosConfig

model = OpenMythos(MythosConfig()).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

input_ids = torch.randint(0, 32000, (2, 512)).cuda()
labels    = torch.randint(0, 32000, (2, 512)).cuda()

logits = model(input_ids)                    # (2, 512, 32000)
loss   = torch.nn.functional.cross_entropy(
    logits.view(-1, 32000),
    labels.view(-1),
)
loss.backward()
optimizer.step()

Depth extrapolation at inference

A looped transformer trained on N loops can be evaluated on N + k loops and often achieves higher quality on hard multi-hop problems. Pass n_loops at inference time:

python
# Trained with max_loop_iters=16 โ€” try deeper reasoning at test time
logits_deep = model(input_ids, n_loops=32)

`generate`

python
@torch.no_grad()
def generate(
    self,
    input_ids: torch.Tensor,
    max_new_tokens: int = 64,
    n_loops: int = 8,
    temperature: float = 1.0,
    top_k: int = 50,
) -> torch.Tensor

Autoregressive token generation with KV caching. Processes the full prompt on step 0, then decodes one token at a time using the accumulated cache.

Parameters

Parameter Type Default Description
input_ids Tensor (B, T) โ€” Prompt token indices
max_new_tokens int 64 Number of new tokens to generate
n_loops int 8 Recurrent loop depth per decode step. Can be higher than the training value for harder prompts (depth extrapolation)
temperature float 1.0 Softmax temperature applied to logits before sampling. Values < 1 make the distribution more peaked (less random); values > 1 make it flatter
top_k int 50 Restricts sampling to the top-K most probable tokens at each step. 0 disables filtering (full vocabulary sampling)

Returns

Tensor (B, T + max_new_tokens) โ€” the original prompt concatenated with the generated token indices.

KV caching mechanism

On step 0, the full prompt (B, T) is passed and all keys/values for every layer are populated in kv_cache. On steps 1โ€ฆN only the single most recent token (B, 1) is passed; the attention layers read back all prior K/V from the cache. This makes decode cost proportional to a single token per step rather than the full growing sequence.

Each layer caches under a deterministic string key ("prelude_0", "recurrent_loop_3", "coda_1", etc.), so caches from different layers never collide.

Sampling strategy

text
logits = forward(cur_ids, n_loops, kv_cache)[:, -1, :] / temperature

if top_k > 0:
    threshold = logits.topk(top_k).values[:, -1:]
    logits[logits < threshold] = -inf

probs    = softmax(logits)
next_tok = multinomial(probs, num_samples=1)

Generation example

python
import torch
from open_mythos.main import OpenMythos, MythosConfig

model = OpenMythos(MythosConfig()).eval()

# Tokenized prompt (use your tokenizer of choice)
prompt = torch.tensor([[1, 450, 3118, 310, 278]])   # (1, 5)

output = model.generate(
    prompt,
    max_new_tokens=128,
    n_loops=16,        # deeper reasoning
    temperature=0.8,
    top_k=40,
)
# output.shape == (1, 133)

Internal Components

The following sub-modules are assembled inside OpenMythos. They are not typically called directly but understanding them clarifies the model's behavior.

`RecurrentBlock`

The heart of the architecture. A single TransformerBlock (with MoE FFN) is run in a loop for up to n_loops iterations, with the following per-iteration pipeline:

text
h_loop = loop_index_embedding(h, t, loop_dim)   # inject sinusoidal loop-index signal
combined = RMSNorm(h_loop + e)                   # add frozen encoded input
trans_out = TransformerBlock(combined, ...)       # attention + MoE FFN
trans_out = trans_out + LoRAAdapter(trans_out, t) # depth-wise LoRA delta
h = LTIInjection(h, e, trans_out)               # stable update: Aยทh + Bยทe + trans_out
p = ACTHalting(h)                                # per-position halting probability

The loop exits early for positions whose cumulative halting probability exceeds cfg.act_threshold. If all positions have halted, the loop exits before n_loops. The final output is an ACT-weighted sum of h across iterations.

`LTIInjection`

Implements the stable recurrent update rule h_{t+1} = Aยทh_t + Bยทe + transformer_out. The diagonal matrix A is parameterized as:

text
A_continuous = Diag(-exp(log_A))     # always negative diagonal
A_discrete   = exp(ฮ”t ยท A_continuous) # ZOH discretization, values โˆˆ (0, 1)

This guarantees spectral radius ฯ(A) < 1 by construction, making the looped model unconditionally stable regardless of learning rate or batch noise. See Parcae (Prairie et al., 2026) for the theoretical foundation.

`ACTHalting`

A single linear layer mapping (B, T, dim) โ†’ (B, T) followed by sigmoid. At each loop step, the scalar halting probability per position is accumulated. When the cumulative sum exceeds cfg.act_threshold, the ACT remainder trick assigns the remaining probability mass as the final weight and the position stops contributing. Implements Graves (2016) ACT.

`LoRAAdapter`

A depth-wise low-rank adapter with three components:

  • down: shared Linear(dim, rank) โ€” down-projects the transformer output
  • B: shared parameter matrix (rank, dim) โ€” up-projects back to full dimension
  • scale: Embedding(max_loops, rank) โ€” per-loop element-wise scale

The delta per iteration is (down(x) * scale[t]) @ B. Bridges the expressiveness gap between pure weight-tying and fully distinct per-layer weights. Based on Relaxed Recursive Transformers (Bae et al., 2024).

`TransformerBlock`

Pre-norm transformer block with swappable attention and FFN:

  • Attention: MLAttention (MLA) or GQAttention (GQA), selected by cfg.attn_type
  • FFN: MoEFFN (when use_moe=True, inside RecurrentBlock) or dense Expert (Prelude, Coda)
  • Pre-norm via RMSNorm applied to both the attention input and FFN input

`MLAttention`

Multi-Latent Attention (DeepSeek-V2, 2024). The cache stores only the compressed KV latent c_kv (rank kv_lora_rank) plus the RoPE-encoded keys. At each decode step, K_nope and V are cheaply reconstructed from c_kv via a shared up-projection, trading a fast linear multiply for dramatically smaller KV memory footprint.

Cache size per layer per token: kv_lora_rank + n_heads ร— qk_rope_head_dim vs. full GQA cache of n_kv_heads ร— head_dim ร— 2.

`GQAttention`

Grouped Query Attention (Ainslie et al., 2023). n_kv_heads KV pairs are shared across n_heads // n_kv_heads query heads each, reducing KV cache by that factor while preserving full query expressiveness.

`MoEFFN`

Fine-grained Mixture-of-Experts FFN (DeepSeekMoE, Dai et al., 2024):

  • Routed experts: n_experts small SwiGLU FFNs. Each token's router selects the top-n_experts_per_tok via softmax over learned logits. A per-expert bias router_bias (non-gradient, updated externally) keeps load balanced.
  • Shared experts: n_shared_experts always-active FFNs with width expert_dim ร— n_experts_per_tok, absorbing cross-domain patterns.

Total activated parameters per token: (n_experts_per_tok / n_experts) of routed capacity + all shared capacity.

`Expert`

Single SwiGLU feed-forward unit: down(silu(gate(x)) * up(x)). Used both as individual routed experts inside MoEFFN and as the dense FFN in Prelude/Coda blocks.

`RMSNorm`

Root Mean Square Layer Normalization (Zhang & Sennrich, 2019). Normalizes by x / rms(x) with a learned per-channel rescaling weight. No bias, no mean subtraction. Used throughout in place of standard LayerNorm.


Utility functions

`precompute_rope_freqs(dim, max_len, theta)`

Precomputes complex-valued RoPE rotation matrices as a (max_len, dim//2) complex64 tensor. Called once in __init__ and stored as a buffer.

`apply_rope(x, freqs_cis)`

Applies precomputed RoPE frequencies to a query or key tensor by treating adjacent feature pairs as complex numbers and multiplying pointwise by the positional phasor.

`loop_index_embedding(h, loop_t, loop_dim, theta)`

Injects a sinusoidal loop-index signal into the first loop_dim channels of the hidden state, analogous to RoPE but over recurrence depth rather than sequence position. Allows the shared recurrent block weights to behave differently at different loop iterations.


Key design properties

Property Mechanism Benefit
Depth extrapolation Recurrent block with looped identical weights Train on N loops, test on N+k โ€” harder problems solved without retraining
Parameter efficiency Weight sharing across all loop iterations k-layer model achieves quality of kL-layer model; parameters โ‰ˆ k, compute โˆ L
Adaptive compute ACT halting per position Easy tokens exit early; hard tokens receive full loop depth โ€” within the same batch
Stable training LTI injection with ZOH-constrained A (ฯ(A) < 1) No residual explosion; robust to high learning rates
Domain breadth MoE FFN in recurrent block Different expert subsets can be routed to at each loop depth
Loop differentiation Loop-index sinusoidal embedding Same weights implement functionally distinct phases per iteration
Efficient KV memory MLA (default) or GQA MLA: 10โ€“20ร— smaller cache vs standard attention at production scale
Depth-wise adaptation LoRA adapter per loop iteration Expressiveness beyond pure weight-tying; minimal parameter overhead

Full configuration reference

The default MythosConfig() targets a mid-scale research model. Below is a minimal configuration for quick experimentation:

python
from open_mythos.main import OpenMythos, MythosConfig

# Minimal config for fast iteration / unit testing
small_cfg = MythosConfig(
    vocab_size=8192,
    dim=256,
    n_heads=4,
    n_kv_heads=2,
    max_seq_len=512,
    max_loop_iters=4,
    prelude_layers=1,
    coda_layers=1,
    attn_type="gqa",
    n_experts=8,
    n_shared_experts=1,
    n_experts_per_tok=2,
    expert_dim=64,
    lora_rank=4,
)
model = OpenMythos(small_cfg)

And a production-oriented MLA configuration matching the default hyperparameters:

python
# Default MLA config (matches MythosConfig() defaults)
prod_cfg = MythosConfig(
    vocab_size=32000,
    dim=2048,
    n_heads=16,
    n_kv_heads=4,
    max_seq_len=4096,
    max_loop_iters=16,
    prelude_layers=2,
    coda_layers=2,
    attn_type="mla",           # Multi-Latent Attention
    kv_lora_rank=512,
    q_lora_rank=1536,
    qk_rope_head_dim=64,
    qk_nope_head_dim=128,
    v_head_dim=128,
    n_experts=64,
    n_shared_experts=2,
    n_experts_per_tok=4,
    expert_dim=512,
    act_threshold=0.99,
    rope_theta=500000.0,
    lora_rank=16,
)
model = OpenMythos(prod_cfg)

References

Component Paper
Recurrent-Depth Transformer Loop, Think, & Generalize (2025)
LTI-stable injection (Parcae) Scaling Laws for Stable Looped Language Models (Prairie et al., 2026)
Looped transformer reasoning Reasoning with Latent Thoughts (Saunshi et al., 2025)
Multi-Latent Attention DeepSeek-V2 (2024)
Grouped Query Attention Ainslie et al., 2023
Mixture-of-Experts FFN DeepSeekMoE (Dai et al., 2024)
Adaptive Computation Time Graves, 2016
Depth-wise LoRA Relaxed Recursive Transformers (Bae et al., 2024)
RMSNorm Zhang & Sennrich, 2019
RoPE Su et al., 2021
Universal Transformer (ACT basis) Dehghani et al., 2018
Continuous latent reasoning COCONUT (2024)

File: training/requirements.txt

torch>=2.11.0
datasets>=3.6.0
loguru>=0.7.3
open-mythos