DOCSEARCHbyAlgolia
Submit Docs

OpenMythos

repository·main·Indexed 12 days ago

https://github.com/kyegomez/openmythos

An open-source theoretical reconstruction of the Claude Mythos Recurrent-Depth Transformer (RDT) architecture. Version 0.5.0 implements depth-variable reasoning through looped layers, featuring support for Multi-Latent Attention (MLA), Grouped Query Attention (GQA) with Flash Attention 2, and Mixture of Experts (MoE). It includes pre-configured model variants from 1B to 1T parameters and tools for training stability via LTI constraints.

Tokens
11.8K
Snippets
31
Records
51
Agent score
97%

What's inside OpenMythos

  1. How Adaptive Computation Time (ACT) works in OpenMythos

    main

    OpenMythos uses ACTHalting to implement adaptive compute. At each loop iteration, a linear layer calculates a scalar halting probability per position.

    As the loop progresses, these probabilities are accumulated. When the cumulative sum exceeds the cfg.act_threshold, the position stops contributing to the loop. The final output is an ACT-weighted sum of the hidden state h across all iterations, allowing easy tokens to exit early while hard tokens receive full loop depth.

  2. How Mixture-of-Experts (MoE) is implemented in the FFN

    main

    The MoEFFN provides fine-grained routing within the recurrent block. It consists of two types of experts:

    1. Routed experts: n_experts small SwiGLU FFNs. A router selects the top-n_experts_per_tok experts for each token via a softmax over learned logits. A router_bias is used to maintain load balance.
    2. Shared experts: n_shared_experts that are always active, absorbing cross-domain patterns.

    This allows the model to use different expert subsets at different loop depths, increasing domain breadth.

  3. Differentiate loop iterations using loop index embeddings

    main
    To prevent the looped block from behaving identically on every iteration, you can inject a loop index embedding (similar to RoPE) alongside the input at each step. This allows the same shared weights to implement functionally distinct operations (e.g., early-stage pattern matching vs. late-stage refinement) across different loop depths.

  4. Prevent 'overthinking' using Adaptive Computation Time (ACT)

    main
    Excessive recurrence can cause the hidden state to drift into noise, a failure mode known as "overthinking." To mitigate this, implement an Adaptive Computation Time (ACT) halting mechanism. This uses a learned scalar per position to dynamically decide when to stop looping, allowing harder tokens more computation and simple tokens to halt early.

  5. How the RecurrentBlock and LTIInjection work together

    main

    The RecurrentBlock is the core of the architecture, running a TransformerBlock in a loop for up to n_loops iterations. To ensure training stability and prevent residual explosion, it uses LTIInjection for the recurrent update.

    LTIInjection implements the rule h_{t+1} = A·h_t + B·e + transformer_out. The diagonal matrix A is parameterized using ZOH (Zero-Order Hold) discretization: A_discrete = exp(Δt · A_continuous), where A_continuous is always a negative diagonal. This guarantees a spectral radius ρ(A) < 1, making the model unconditionally stable regardless of learning rate or batch noise.

  6. Balance reasoning and memorization with looping-based regularization

    main
    Looped architectures are structurally biased toward composition (reasoning) over memorization (rote facts). To balance this tradeoff during training, use looping-based regularization: apply stronger looping constraints for reasoning tasks and relax them for retrieval/memorization tasks.

  7. Ensure training stability in looped models via LTI constraints

    main

    Training looped models is prone to residual explosion (unbounded growth of hidden state h_t) and loss spikes. To solve this, treat the recurrence as a discrete linear time-invariant (LTI) dynamical system: h_{t+1} = A·h_t + B·e.

    Stability is guaranteed if the spectral radius of A is less than 1 (ρ(A) < 1). You can enforce this by construction using the Parcae architecture approach:

    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.

    This ensures the model remains stable regardless of learning rate or batch noise.

    h_{t+1} = A·h_t + B·e

  8. Understand scaling laws for looped models

    main

    Looped models follow predictable scaling laws that differ from fixed-depth Transformers:

    • Training Scaling: For a fixed FLOP budget and fixed parameters, increasing mean recurrence while reducing token count yields lower loss than training with minimal loops on more data. Both optimal recurrence and optimal token count follow power laws.
    • Inference Scaling: Increasing test-time loops improves quality following a predictable, saturating exponential decay. This is similar to the inference-time scaling seen in chain-of-thought reasoning.

    Efficiency Note: A looped model at 770M parameters can achieve the downstream quality of a 1.3B fixed-depth Transformer, effectively providing similar quality with roughly half the parameters.

  9. How the Recurrent-Depth Transformer (RDT) architecture works

    main

    OpenMythos implements a Recurrent-Depth Transformer (RDT), also known as a Looped Transformer. The architecture divides layers into three functional blocks:

    1. Prelude (P): Standard transformer layers run once.
    2. Recurrent Block (R): A subset of layers that are looped $T$ times. During each loop, the hidden state $h$ is updated using an input injection $e$ from the Prelude.
    3. Coda (C): Standard transformer layers run once after the loops.

    The recurrent update rule at each loop step $t$ is:

    $$h_{t+1} = A \cdot h_t + B \cdot e + \text{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 to prevent signal drift.
    • $A$ and $B$ are learned injection parameters.
    • $\text{Transformer}$ applies attention and MLP blocks.

    This architecture allows for deeper reasoning (more loops) without increasing the parameter count.

  10. Optimize inference with Continuous Depth-wise Batching

    main
    Because tokens in a recursive architecture share the same recurrent block, you can implement Continuous Depth-wise Batching. This allows the model to exit the loop at different depths for different tokens or sequences within the same batch, processing easy inputs quickly and hard inputs with more iterations. This can improve inference throughput by an estimated 2-3x.

  11. Adapt loop behavior using depth-wise LoRA

    main
    To allow each loop to adapt its behavior slightly without the overhead of fully distinct layers, you can add a small depth-wise LoRA module at each iteration. This approach (from Relaxed Recursive Transformers) uses a large common weight matrix (the recursive base) and a small rank-r adaptation matrix that shifts behavior per iteration depth.

  12. How LoRAAdapter provides depth-wise adaptation

    main

    To bridge the gap between weight-tying (sharing weights across loops) and having fully distinct per-layer weights, LoRAAdapter is used. It applies a depth-wise low-rank delta at each loop iteration.

    The delta is calculated as (down(x) * scale[t]) @ B, where:

    • down: A shared Linear(dim, rank) projection.
    • B: A shared parameter matrix (rank, dim) for up-projection.
    • scale: An Embedding(max_loops, rank) that provides a per-loop element-wise scale.