mixture_of_recursions: a router that gives every token a depth

Early exiting leaves two problems behind

The research motivation in the README starts from CALM, an early-exiting adaptive computation technique that lets a language model skip unnecessary computation when it is confident in its predictions. That creates dynamic pathways through the model, and the point of the paper is that putting early exiting into practice hits two main bottlenecks.

The first is the missing key-value cache problem. When tokens exit early, they skip computing the KV pairs for the remaining deeper layers, and those missing values are essential for decoding future tokens. Approximating them tends to hurt performance, which is the problem the first bottleneck creates.

The second is inefficient batched inference. Tokens that exit early end up waiting for others in the same batch to finish their full computation, and that idling prevents efficient batching and wastes processing time.

Those two problems are why early exiting is not simply a speed trick. A model that skips work per token saves computation in the single-sequence case and loses most of the benefit in the batched case that production serving actually uses, and it also loses the context that later tokens need. Any method that claims to fix early exiting has to answer both questions, and this repository exists because the authors believe they have.

A learned router picks the depth, and the cache follows the depth

The proposed method is MoR, and its two components map directly onto the two bottlenecks.

The first is a routing mechanism trained end to end, which dynamically assigns the optimal recursion depth to each token. Instead of a fixed rule that decides when a token is confident enough to stop, the depth is a learned decision, and the training signal comes from the same objective as the rest of the model. That is what makes the depth assignment per token rather than per sequence.

The second is a recursion-wise KV caching strategy that selectively stores KV pairs. Since the router produces tokens at different depths, the cache is organised by recursion level rather than by layer, which is what resolves the missing cache problem while also optimising memory usage.

The reported result is up to 2x greater inference throughput compared to standard transformers at similar accuracy, along with reduced total training FLOPs and memory requirements. The wording of that claim matters for a reader: it is an upper bound on a specific comparison rather than a general result, and the paper rather than the repository is where the evaluation conditions live.

The implementation overview section of the README is cut off after its first heading, so the detail on dataset handling, the model code and the configuration system has to come from the paper and the repository itself.

How this differs from FREE and Recursive Transformers

The README positions the work against two pieces of prior art, and the comparison is specific enough to be useful.

The FREE Framework addressed the missing KV cache problem using parallel decoding, described as an early form of self-speculative decoding, which efficiently computes the exact KV pairs for early-exited tokens. The limitation the README names is that compatibility with batch inference remained poor, so it solved the cache problem and left the idling problem in place.

Recursive Transformers aimed to mitigate inefficient batched inference through parameter sharing, enabling tokens at different depths to be processed together. The limitation named here is different in kind: two separate training processes for integrating parameter sharing and early exiting degraded performance, and the model still required handling the missing KV cache.

So the two predecessors split the problem in half, one fixing the cache and one fixing the batching, each at a cost. MoR is described as a unified framework that directly tackles both, with the router trained end to end.

For a reader choosing a starting point, that is the substance of the contribution. The price is a more complex training setup, since a learned router has to be trained together with the model rather than bolted on afterwards.

Building the environment, and the versions the authors used

The setup instructions are conda based, and the README recommends conda for dependency management.

The environment is created and activated like this:

    conda create -n mor python=3.12
    conda activate mor

Then the packages are installed, with the recommendation to bring pip and setuptools up to date first:

    pip install --upgrade pip
    pip install --upgrade setuptools
    pip install torch
    pip install -r requirements.txt
    # If you experience issues with flash-attn, try:
    # pip install flash-attn --no-build-isolation

The comment about flash-attn is not decoration. FlashAttention compiles from source and often fails under build isolation, and the README gives the exact fallback command for that case.

The version guidance is explicit as well. The authors state that they specifically used torch==2.6.0+cu124, flash_attn==2.7.4.post1 and transformers==4.52.4, and that if you encounter issues you should consider these exact versions. That is a research reproducibility note rather than a minimum requirement, but in practice the pinned combination is the one that is known to work with the code, and the CUDA build suffix indicates a specific GPU toolchain.

The repository itself has no GitHub releases, so there is no packaged environment to download and the requirements file is the only installation contract.

FineWeb-Edu through directories and symbolic links

The models are pretrained on a deduplicated subset of the FineWeb-Edu dataset, which is available as part of the SmolLM-Corpus on Hugging Face. The preparation steps are laid out in detail because the layout matters to the scripts.

First you create the directory structure under your own data path, replacing the placeholder with your real path:

    mkdir -p {your_data_path}/mixture_of_recursions/hf_cache
    mkdir -p {your_data_path}/mixture_of_recursions/hf_datasets
    mkdir -p {your_data_path}/mixture_of_recursions/hf_models
    mkdir -p {your_data_path}/mixture_of_recursions/results

Then you link that data path into the project, with the project path as the second placeholder:

    ln -s {your_data_path}/mixture_of_recursions/* {your_project_path}/mixture_of_recursions/

The corpus itself is fetched with the provided script:

    bash lm_dataset/download_scripts/download_fineweb-edu-dedup.sh

Finally the cached dataset has to be moved to the DATA_DIR that the loader expects, and the README points at the specific lines of lm_dataset/load_dataset.py that define it, noting that steps three and four can be done together using the script download_langauge_modeling_datasets.sh, spelled with that typo in the filename.

The symlink arrangement is there so that the corpus, the model checkpoints and the results live outside the source tree. It keeps a large download out of the working copy, at the cost of a setup step that has to be done in the right order.

The requirements file describes the whole stack

The dependency list is worth reading as a description of what the project actually does, because it separates training from evaluation.

For training the file lists torch, hydra-core, omegaconf, wandb, tensorboard, transformers==4.52.4, datasets, accelerate, deepspeed, peft, zstandard, boto3, smart_open, bitsandbytes and flash-attn. Hydra and OmegaConf are the configuration system, which is confirmed by the conf/ and ds_configs/ directories in the repository. wandb and tensorboard are the experiment trackers, deepspeed and bitsandbytes are the distributed training and mixed precision pieces, peft is there for parameter efficient fine tuning, and zstandard with smart_open and boto3 suggests the data pipeline can read compressed data from object storage.

For evaluation there is a separate group: sacrebleu, evaluate, sqlitedict, nltk, pytablewriter, rouge_score, matplotlib and more_itertools, which covers translation metrics, model evaluation harnesses, Rouge scoring and result tables.

The last group is the language model evaluation harness, and the instructions are to clone it from the EleutherAI repository and install it in editable mode:

    git clone --depth 1 https://github.com/EleutherAI/lm-evaluation-harness
    cd lm-evaluation-harness
    pip install -e .

One subtlety worth flagging: the flash-attn comment in the file matches the workaround in the setup instructions, which is a reminder that requirements files for research code usually carry known workarounds rather than pure specifications.

A research pipeline, not a package

The top level of the repository shows what you get if you clone it: acc_configs/ and ds_configs/ for the configuration trees, conf/ for the main configuration, lm_dataset/ for the data pipeline, model/ for the model code, lm_eval/ for the evaluation code, util/ and scripts/ for supporting code, notebooks/ for exploration, and asset/ for figures. The entry points at the root are pretrain.py, eval_fewshot.py and evaluate_fineweb_test.py, plus paths.py for path resolution.

That layout describes a complete research pipeline: pretrain, evaluate few-shot tasks, and evaluate on the fineweb test set. It is not a library, there is no package entry point, and there is nothing to install for inference use.

The licence is Apache-2.0, which is permissive for research code and for derived work, with the usual patent grant. The last push was on 2026-09-19 and the repository has no GitHub releases, so there are no versioned artefacts and the paper, not the repository, is the citation of record. The paper is on arXiv at 2507.10524 and the authors list covers KAIST AI, Mila, Google Cloud, Google DeepMind, Google Research and Université de Montréal, with the paper badge in the README pointing at the arXiv entry and a BibTeX badge for citation.

For anyone reproducing the work, the README's own advice is the practical summary: match the exact versions, prepare the data directories before the download, and expect to read the paper for the method details that the README does not spell out.

Editorial conclusion

Adopt mixture_of_recursions if you are doing research on adaptive computation and want to reproduce or extend the Mixture-of-Recursions setup, with a NVIDIA GPU, a CUDA stack and the patience to download a pretraining corpus. Skip it if you want a production inference server, since the repository is a research pipeline with a training script, evaluation scripts and configuration directories rather than a packaged runtime. Before you start, pin the versions the README names, torch 2.6.0 with cu124, flash_attn 2.7.4.post1 and transformers 4.52.4, and expect to read the arXiv paper for the parts of the implementation overview that the README leaves out. The last push was on 2026-09-19, the licence is Apache-2.0, and there are no GitHub releases.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. Project website
4. raymin0223/mixture_of_recursions on GitHub
5. README