openmythos-recurrent-transformer

What this skill does

```markdown
---
name: openmythos-recurrent-transformer
description: Build and experiment with Recurrent-Depth Transformer (RDT) models using OpenMythos, a theoretical reconstruction of the Claude Mythos architecture with looped transformers, MLA/GQA attention, and sparse MoE.
triggers:
  - implement a looped transformer model
  - build a recurrent depth transformer
  - use OpenMythos for inference time scaling
  - configure MLA or GQA attention in OpenMythos
  - set up mixture of experts with recurrent blocks
  - train a model with adaptive loop iterations
  - generate text with variable reasoning depth
  - explore compute-adaptive transformer architectures
---

# OpenMythos Recurrent-Depth Transformer

> Skill by [ara.so](https://ara.so) — Daily 2026 Skills collection.

OpenMythos is an open-source theoretical implementation of a Recurrent-Depth Transformer (RDT) inspired by the Claude Mythos architecture. It divides computation into three stages: **Prelude** (standard transformer layers run once), a **Recurrent Block** (looped up to `max_loop_iters` times with stable injection), and a **Coda** (standard transformer layers run once). Attention is switchable between MLA (Multi-head Latent Attention) and GQA (Grouped Query Attention), and feed-forward layers use sparse MoE with routed and shared experts.

## Installation

```bash
git clone https://github.com/The-Swarm-Corporation/OpenMythos.git
cd OpenMythos
pip install -r requirements.txt
```

## Core Concepts

### Architecture Flow
```
Input → Embedding
  ↓
[Prelude]          — N standard transformer layers, run once
  ↓
[Recurrent Block]  — looped T times per forward pass
  ↑_________↓        h_{t+1} = A·h_t + B·e + Transformer(h_t, e)
  ↓
[Coda]             — N standard transformer layers, run once
  ↓
Output Logits
```

- `e` = encoded input from Prelude (injected every loop to prevent drift)
- `A`, `B` = learned stable injection parameters (spectral radius ρ(A) < 1 enforced)
- More loops at inference = deeper implicit reasoning (no extra parameters)

### Attention Types
- **MLA** (`"mla"`): Multi-head Latent Attention — uses KV LoRA compression, separate RoPE and NoPE head dimensions
- **GQA** (`"gqa"`): Grouped Query Attention — fewer KV heads than Q heads, simpler config

## Configuration Reference

### `MythosConfig` Parameters

| Parameter | Type | Description |
|---|---|---|
| `vocab_size` | int | Vocabulary size |
| `dim` | int | Model hidden dimension |
| `n_heads` | int | Number of attention query heads |
| `n_kv_heads` | int | Number of KV heads (GQA ratio = n_heads/n_kv_heads) |
| `max_seq_len` | int | Maximum sequence length |
| `max_loop_iters` | int | Maximum recurrent loop iterations |
| `prelude_layers` | int | Number of Prelude transformer layers |
| `coda_layers` | int | Number of Coda transformer layers |
| `n_experts` | int | Total number of MoE routed experts |
| `n_shared_experts` | int | Always-active shared experts |
| `n_experts_per_tok` | int | Top-K experts selected per token |
| `expert_dim` | int | Hidden dim inside each expert FFN |
| `lora_rank` | int | LoRA rank for injection parameters |
| `attn_type` | str | `"mla"` or `"gqa"` |
| `kv_lora_rank` | int | MLA only: KV compression rank |
| `q_lora_rank` | int | MLA only: Q compression rank |
| `qk_rope_head_dim` | int | MLA only: RoPE head dimension |
| `qk_nope_head_dim` | int | MLA only: NoPE head dimension |
| `v_head_dim` | int | MLA only: Value head dimension |

## Usage Patterns

### Pattern 1: GQA Model (Simpler Config)

```python
import torch
from open_mythos.main import OpenMythos, MythosConfig

cfg = MythosConfig(
    vocab_size=32000,
    dim=512,
    n_heads=8,
    n_kv_heads=2,          # GQA: 4x fewer KV heads
    max_seq_len=2048,
    max_loop_iters=8,
    prelude_layers=2,
    coda_layers=2,
    n_experts=16,
    n_shared_experts=2,
    n_experts_per_tok=2,
    expert_dim=256,
    lora_rank=16,
    attn_type="gqa",
)

model = OpenMythos(cfg)
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

# Forward pass with 4 reasoning loops
ids = torch.randint(0, cfg.vocab_size, (2, 64))
logits = model(ids, n_loops=4)
print(f"Logits shape: {logits.shape}")  # (2, 64, 32000)
```

### Pattern 2: MLA Model (More Expressive)

```python
import torch
from open_mythos.main import OpenMythos, MythosConfig

cfg = MythosConfig(
    vocab_size=32000,
    dim=512,
    n_heads=8,
    n_kv_heads=8,           # MLA uses full heads
    max_seq_len=2048,
    max_loop_iters=16,
    prelude_layers=2,
    coda_layers=2,
    n_experts=16,
    n_shared_experts=2,
    n_experts_per_tok=2,
    expert_dim=256,
    lora_rank=16,
    attn_type="mla",
    # MLA-specific compression dims
    kv_lora_rank=64,
    q_lora_rank=128,
    qk_rope_head_dim=32,
    qk_nope_head_dim=32,
    v_head_dim=32,
)

model = OpenMythos(cfg)

ids = torch.randint(0, cfg.vocab_size, (1, 128))

# More loops = deeper implicit reasoning
logits_shallow = model(ids, n_loops=2)
logits_deep    = model(ids, n_loops=16)
print(f"Shallow logits: {logits_shallow.shape}")
print(f"Deep logits:    {logits_deep.shape}")
```

### Pattern 3: Text Generation with Variable Depth

```python
import torch
from open_mythos.main import OpenMythos, MythosConfig

cfg = MythosConfig(
    vocab_size=50257,
    dim=768,
    n_heads=12,
    n_kv_heads=4,
    max_seq_len=1024,
    max_loop_iters=12,
    prelude_layers=2,
    coda_layers=2,
    n_experts=8,
    n_shared_experts=1,
    n_experts_per_tok=2,
    expert_dim=512,
    lora_rank=32,
    attn_type="gqa",
)

model = OpenMythos(cfg)
model.eval()

prompt_ids = torch.randint(0, cfg.vocab_size, (1, 16))

# Simple task: fewer loops
with torch.no_grad():
    easy_output = model.generate(
        prompt_ids,
        max_new_tokens=32,
        n_loops=2,          # shallow reasoning
    )

# Complex task: more loops (same model, more compute)
with torch.no_grad():
    hard_output = model.generate(
        prompt_ids,
        max_new_tokens=32,
        n_loops=12,         # deep reasoning
    )

print(f"Easy output shape: {easy_output.shape}")
print(f"Hard output shape: {hard_output.shape}")
```

### Pattern 4: Stability Verification

Always verify the spectral radius constraint after initialization and training:

```python
import torch
from open_mythos.main import OpenMythos, MythosConfig

cfg = MythosConfig(
    vocab_size=1000,
    dim=256,
    n_heads=8,
    n_kv_heads=2,
    max_seq_len=128,
    max_loop_iters=8,
    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="gqa",
)

model = OpenMythos(cfg)

# Check injection matrix spectral radius — must be < 1
A = model.recurrent.injection.get_A()
spectral_radius = A.max().item()
print(f"Spectral radius ρ(A): {spectral_radius:.6f}")
assert spectral_radius < 1.0, "UNSTABLE: spectral radius >= 1!"
print("✓ Stability constraint satisfied")
```

### Pattern 5: Scaling Experiment — Loops vs Quality

```python
import torch
import torch.nn.functional as F
from open_mythos.main import OpenMythos, MythosConfig

cfg = MythosConfig(
    vocab_size=1000,
    dim=256,
    n_heads=8,
    n_kv_heads=2,
    max_seq_len=64,
    max_loop_iters=16,
    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="gqa",
)

model = OpenMythos(cfg)
model.eval()

ids = torch.randint(0, cfg.vocab_size, (1, 32))
targets = torch.randint(0, cfg.vocab_size, (1, 32))

results = {}
with torch.no_grad():
    for n_loops in [1, 2, 4, 8, 16]:
        logits = model(ids, n_loops=n_loops)
        loss = F.cross_entropy(
            logits.view(-1, cfg.vocab_size),
            targets.view(-1)
        )
        results[n_loops] = loss.item()
        print(f"n_loops={n_loops:2d} → loss={loss.item():.4f}")

# Expect diminishing returns (saturating exponential decay)
print("\nDelta losses (should decrease):")
loop_counts