Ouro-1.4B

📚 Paper • 🏠 Project Page

Fork Description

• Makes the Ouro model compatible with Transformers v5 (tested with 5.9.0).
• Fixes several smaller bugs.
• Note: Proper early exit will be implemented soon. The upstream repo currently only supports post-hoc gating, where KVs participate in attention even though they would not exist under real early exit.

Model Description

⚠️ IMPORTANT: This model is intended for research purposes only. It is provided as-is without warranties for production use.

Ouro-1.4B is a 1.4 billion parameter Looped Language Model (LoopLM) that achieves exceptional parameter efficiency through iterative shared-weight computation.

Key Features

• Exceptional Parameter Efficiency: Matches 3-4B standard transformer performance with only 1.4B parameters
• Iterative Latent Reasoning: Performs reasoning through recurrent computation in latent space
• Adaptive Computation: Supports early exit mechanisms for dynamic compute allocation

Configuration

Recurrent Steps and Adaptive Exit

The model's computational behavior can be configured through the config.json file:

{
  "total_ut_steps": 4,
  "early_exit_threshold": 1.0
}

• total_ut_steps: Controls the number of recurrent steps (default: 4). You can adjust this value to trade off between performance and computation time.
• early_exit_threshold: Controls the adaptive exit mechanism (default: 1.0). Lower values encourage earlier exit, while 1.0 means always use all steps.

Example: Modify recurrent steps

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("ByteDance/Ouro-1.4B")
config.total_ut_steps = 3  # Use 3 recurrent steps instead of 4
model = AutoModelForCausalLM.from_pretrained(
    "ByteDance/Ouro-1.4B",
    config=config,
    device_map="auto"
)

Note: vLLM does not currently support the adaptive exit feature due to its inference optimization characteristics. When using vLLM, the model will always execute the full number of total_ut_steps.

Model Architecture

Ouro-1.4B is based on the decoder-only Transformer architecture with parameter sharing across recurrent steps:

Training Details

• Training Tokens: 7.7T tokens
• Training Pipeline:
  • Stage 1: Pre-training (6T tokens)
  • Stage 2: CT Annealing (1.4T tokens)
  • Stage 3: Long Context Training (20B tokens)
  • Stage 4: Mid-training (300B tokens)
• Data Composition: Web data, code, mathematics, long-context documents
• Optimizer: AdamW (β₁=0.9, β₂=0.95)
• Learning Rate Scheduler: Warmup-Stable-Decay (WSD)

Quick Start

⚠️ IMPORTANT: Works with transformers==5.9.0.

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "ByteDance/Ouro-1.4B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype="auto"
)

# Generate text
inputs = tokenizer("The future of AI is", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Acknowledgments

We thank @Antizana for the KV cache fix merged from ouro-cache-fix, which resolved a critical compatibility issue with transformers>=4.56.0.

Citation

@article{zhu2025scaling,
  title={Scaling Latent Reasoning via Looped Language Models},
  author={Zhu, Rui-Jie and Wang, Zixuan and Hua, Kai and Zhang, Tianyu and Li, Ziniu and Que, Haoran and Wei, Boyi and Wen, Zixin and Yin, Fan and Xing, He and others},
  journal={arXiv preprint arXiv:2510.25741},
  year={2025}
}

## License

This model is licensed under Apache-2.0. See the LICENSE file for details.

---
## Trained lookahead exit gate

This checkpoint includes the original `model.early_exit_gate` and a trained `model.lookahead_exit_gate`.
The lookahead head is distilled from the frozen original gate on Nemotron-Post-Training-Dataset-v2 and predicts the next recurrent step's exit hazard.
The model configuration defaults to `exit_gate_type="lookahead"`; set it to `"early_exit"` to select the original head.
Lookahead hazards accumulate into a cumulative exit probability, and a decision computed after depth `d` applies after depth `d + 1`.
Loop-SGLang selects the head through `--json-model-override-args '{"exit_gate_type":"lookahead"}'` and enables early exits only when `--exit-threshold` is supplied.
This gate does not address accuracy loss from shared KV caches.

The lookahead weights come from `outputs/ablations/ouro_gate/lookahead_exit_gate.safetensors` in the continuous-depth-batching research implementation.
Their SHA-256 is `76be752c7744c3b29bbe319a36131814fbe93405e1801ae57c1ba543c735a90c`.