R³-ViT: Recursive, Relaxed, Routed Vision Transformer

Repository: https://github.com/estebancarlin/MoNE-MoR/tree/feat/pretrained

R³-ViT (Recursive, Relaxed, Routed) is a Vision-Transformer extension of the Mixture-of-Recursions method (Bae et al., NeurIPS 2025). A standard ViT runs every token through a fixed stack of L distinct layers. R³-ViT instead folds the stack into a small set of weight-shared stages that are re-applied recursively, and adds a router that decides, per token, how many recursion passes each token receives. Tokens that are "easy" exit early; "hard" tokens recurse deeper. This buys adaptive, token-level compute at a fraction of the parameters.

This repository implements three routing modes plus several efficiency mechanisms:

Origin. This work started from the official Mixture-of-Recursions codebase, which is a language-model project (FineWeb-Edu / SmolLM, the lm_eval harness). That original LM code is kept here for reference but is not the active research direction — the vision side documented below was built on top of it. See Attribution & Citation at the bottom.

---

Documentation map

This top-level README covers the directory layout, a user guide, system configuration, and how to change hyperparameters. Per-function documentation lives in a README inside each major directory:

New to the repo? This is a large fork with a lot of legacy, external, and exploratory code. DEAD_CODE.md is an import-reachability map showing that only 99 of ~1340 .py files are used by the standard train + eval flow — read it first to find the minimal working set and avoid the LM-side / archived directories.

---

1. Directory structure

mixture_of_recursions/
│
├── pretrain.py                # Single Hydra entry point for ALL training (vision + LM)
├── eval_fewshot.py            # Hydra entry point for evaluation (vision + LM)
├── paths.py                   # Resolves PROJECT_ROOT / SAVE_DIR / HF_CACHE_DIR / DATA_DIR / MODEL_DIR
│
├── drun                       # One-line background training launcher (nohup) -> ./drun <config> [num_gpus]
├── run_docker_training.sh     # Full training launcher with GPU-id / delayed-start control
├── eval_and_analyse.sh        # End-to-end per-model eval + routing analysis pipeline
├── how_to_run_on_this_server.txt   # Host script to (re)create the mor_workspace Docker container
│
├── conf/                      # *** Source of truth for every experiment (Hydra + OmegaConf) ***
│   ├── pretrain/              #   Training configs: {c10|c100|i1k}_{poc|pre}_{tc|ec|act|...}_v{N}.yaml
│   └── eval_fewshot/          #   Evaluation configs
│
├── model/                     # Core model implementations
│   ├── util.py                #   get_model / load_model_from_config factory + dtype/config helpers
│   ├── vit_model/             #   *** R³-ViT core: routing transforms + TC/EC/ACT block wrappers ***
│   ├── sharing_strategy/      #   The "fold": average/cycle/middle_cycle weight-sharing (vit.py)
│   ├── relaxation/            #   Per-recursion adapters: SVD-LoRA (lora/), prompt, conditioning
│   ├── distillation/          #   Optional knowledge distillation (logit-KD, dist-token, CKA)
│   ├── mor_model/             #   [legacy LM] Llama-based MoR routers (preserved upstream)
│   ├── base_model/            #   [legacy LM] stock Llama modeling
│   └── recursive_model/       #   [legacy LM] recursive Llama variant
│
├── util/                      # Training & evaluation infrastructure
│   ├── trainer_pt.py          #   MoRTrainer (HF Trainer subclass): optimizer groups, loss aggregation
│   ├── config.py              #   preprocess_config: validates + derives schedules/batch sizes/phases
│   ├── phase_runner.py        #   Two-phase (fold-heal -> A-ViT) orchestration
│   ├── efficiency.py          #   Params / GFLOPs / throughput / latency accounting
│   └── ...                    #   callbacks, losses, metrics, optimizer/scheduler helpers
│
├── vision_dataset/            # Vision data pipeline
│   ├── load_dataset.py        #   CIFAR-10/100 + ImageNet-1k loaders, augmentation factory
│   ├── data_preprocessing.py  #   DeiT-III 3-augment, RandAugment, MixUp, CutMix transforms
│   └── repeat_aug_sampler.py  #   RASampler (DeiT-III repeated augmentation)
│
├── acc_configs/               # Accelerate launcher configs (single_gpu_config.yaml, default_config.yaml)
├── ds_configs/                # DeepSpeed ZeRO configs
├── tests/                     # pytest regression checks (test_*.py); no aggregate runner / CI
├── scripts/                   # Orchestration & analysis shell scripts
├── tools/                     # Checkpoint export/merge utilities
├── notebooks/                 # Analysis notebooks (routing-decision exploration)
│
├── lm_dataset/                # [legacy LM] FineWeb-Edu dataset + tokenization
├── lm_eval/                   # [legacy LM] EleutherAI lm-evaluation-harness (vendored)
│
│   # ---- generated / not under version control (safe to ignore when reading code) ----
├── results/        ⚙ output   # SAVE_DIR: checkpoints, logs, eval artifacts (results/pretrain/<run>/...)
├── logs/           ⚙ output   # Training job logs (logs/vit_<PID>.{out,err})
├── outputs/        ⚙ output   # Hydra run outputs
├── hf_cache/       ⚙ input    # HuggingFace model/dataset cache (large; external data)
├── hf_datasets/    ⚙ input    # Source datasets (DATA_DIR)
├── hf_models/      ⚙ input    # HF model checkpoints for distillation (MODEL_DIR)
│
│   # ---- project docs / analysis (read for research context) ----
├── CLAUDE.md                  # Engineering guide for this repo (architecture, conventions)
├── z_Pretrained/              # Live project journal + pretrained-bootstrap analyses
├── routing_attention_analysis/ # Routing-behaviour & efficiency-benchmark scripts and outputs
└── thesis_assets/             # Figures and bundles for the thesis

Legend: ⚙ input = data the code reads, ⚙ output = artifacts the code writes, [legacy LM] = preserved upstream language-model code (not the active direction).

---

2. User guide

All commands run inside the long-lived Docker container mor_workspace (image mor_environment:ready), whose workdir /home/user1001/mixture_of_recursions is bind-mounted from the host. Host Python does not have the right dependencies. Prefix any Python/CLI command with docker exec mor_workspace ..., or attach a shell with docker exec -it mor_workspace bash. To create a fresh container, run the host launcher in how_to_run_on_this_server.txt.

The unit of work is a config name = the YAML basename in conf/pretrain/ (no .yaml).

2.1 Train

# Simplest: background launch (nohup, logs to logs/vit_<PID>.out), 1 GPU
./drun c100_pre_dense_10rec_svdlora_mixedrank_v1 1

# Full control: choose GPUs and/or a delayed start time
./run_docker_training.sh <config_name> [num_gpus] [gpu_ids] [start_time]
# e.g. 2 GPUs, ids 0 and 3:        ./run_docker_training.sh c100_pre_ec_3rec_v1 2 0,3
# e.g. 1 GPU, start at 18:15:      ./run_docker_training.sh c100_pre_ec_3rec_v1 1 2 18:15

Under the hood both scripts call the Hydra entry point:

docker exec mor_workspace accelerate launch \
  --config_file acc_configs/single_gpu_config.yaml \
  pretrain.py --config-name <config_name>

2.2 Evaluate + routing analysis

# Dataset is inferred from the c10_/c100_/i1k_ prefix of the model name.
# Two-phase (v11+) runs auto-descend into phase1/; override with PHASE=phase0.
bash eval_and_analyse.sh <model_name> [ckpt_step]

# Direct Hydra eval only:
docker exec mor_workspace python eval_fewshot.py --config-name <eval_config_name>

2.3 Quick model-build sanity probe (no training)

Verifies the built model's parameter count without launching a run (the MoR/MoNE transform is applied on top of the dense base, exactly as pretrain.py does it):

docker exec mor_workspace python -c "from model.util import load_model_from_config, normalize_mor_router_config; from model.sharing_strategy import SHARING_STRATEGY; from util.config import preprocess_config; from omegaconf import OmegaConf; cfg = preprocess_config(OmegaConf.load('conf/pretrain/c100_pre_ec_3rec_v1.yaml')); m = load_model_from_config(cfg); m, lora = SHARING_STRATEGY[cfg.model](cfg, m) if cfg.recursive.enable else (m, None); (normalize_mor_router_config(cfg), getattr(m, 'transform_layer_to_mor_'+cfg.mor.type)(cfg, lora_init_dict=lora) if cfg.mor.type=='expert' else getattr(m, 'transform_layer_to_mor_'+cfg.mor.type)(cfg)) if cfg.get('mor',{}).get('enable') else None; print(sum(p.numel() for p in m.parameters())/1e6)"

2.4 Worked example (end to end)

Train the CIFAR-100 dense 10-recursion SVD-LoRA "mixed-rank" model, then evaluate it:

# 1) Train (writes to results/pretrain/c100_pre_dense_10rec_svdlora_mixedrank_v1/)
./drun c100_pre_dense_10rec_svdlora_mixedrank_v1 1

# 2) Watch the log
tail -f logs/vit_*.out

# 3) Evaluate + routing/efficiency analysis once checkpoints exist
bash eval_and_analyse.sh c100_pre_dense_10rec_svdlora_mixedrank_v1

# 4) Results: top-1 accuracy + efficiency report under results/eval_fewshot/final/

2.5 Tests

docker exec mor_workspace pytest tests/test_svd_lora.py -q

---

3. System configuration

All paths are resolved in paths.py relative to PROJECT_ROOT (the repo root). To put large data on a separate disk, create the real directories elsewhere and symlink them in: ln -s {data_disk}/mixture_of_recursions/* {repo}/mixture_of_recursions/.

Output sub-layout under SAVE_DIR:

Launcher configs (which GPUs / distributed backend) live in acc_configs/ (Accelerate: single_gpu_config.yaml, default_config.yaml) and ds_configs/ (DeepSpeed ZeRO). The training scripts pick the single- vs multi-GPU Accelerate config based on the requested GPU count.

---

4. How to modify hyperparameters

Everything is config-driven. A single YAML in conf/pretrain/ defines model architecture, recursion, routing mode, MoNE/MoE, SVD-LoRA, optimizer, schedule, dataset, and augmentation. There are two ways to change a hyperparameter:

(a) Edit / clone a config file (the normal way)

The full per-section reference is in conf/README.md. The highest-impact knobs, with example values taken from conf/pretrain/c100_pre_dense_10rec_svdlora_mixedrank_v1.yaml:

Immutability rule. A config that already has a results/pretrain/<name>/ directory is frozen — never edit it in place. Copy it to a new file with a bumped _vN suffix (or a descriptive tag) and edit the copy. This keeps every result reproducible from its config. (See CLAUDE.md and conf/README.md.)

(b) Override on the command line (Hydra, for quick one-offs)

Any field can be overridden at launch without editing the file:

docker exec mor_workspace python pretrain.py --config-name c100_pre_ec_3rec_v1 \
  learning_rate=5e-4 num_train_steps=30000 recursive.num_recursion=4

util/config.py::preprocess_config then validates the result and derives anything implied (effective batch size, warmup defaults, phase-step consistency).

---

Attribution & Citation

This repository was built starting from raymin0223/mixture_of_recursions (paper: arXiv:2507.10524), the official Mixture-of-Recursions codebase for language models. The vision (ViT) extension documented here is the contribution; the original language-model code, the FineWeb-Edu/SmolLM pipeline, and the vendored lm_eval harness are kept from that base. If you use this work, please cite the original paper:

@misc{bae2025mixtureofrecursionslearningdynamicrecursive,
    title={Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computation},
    author={Sangmin Bae and Yujin Kim and Reza Bayat and Sungnyun Kim and Jiyoun Ha and Tal Schuster and Adam Fisch and Hrayr Harutyunyan and Ziwei Ji and Aaron Courville and Se-Young Yun},
    year={2025},
    eprint={2507.10524},
    archivePrefix={arXiv},
    primaryClass={cs.CL},
    url={https://arxiv.org/abs/2507.10524},
}