File size: 5,541 Bytes
17bc7ca
4634a7d
 
17bc7ca
4634a7d
 
 
 
 
 
 
17bc7ca
4634a7d
ca33b7b
4634a7d
a29f310
 
d509ff8
3cea8f5
d509ff8
4634a7d
a29f310
4634a7d
 
 
 
 
 
 
 
 
 
 
 
3cea8f5
4634a7d
 
43ed691
4634a7d
 
 
 
 
 
 
43ed691
 
4634a7d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
006866a
4634a7d
 
c044851
 
 
 
006866a
 
c044851
7784102
 
3cea8f5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c044851
 
a29f310
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
---
language:
- en
license: apache-2.0
library_name: transformers
pipeline_tag: text-generation
tags:
- arloop
- recurrent
- custom_code
- pretrained
---

# ARLoop-1.2B

📚 [Paper](https://arxiv.org/abs/2611.05234) · 💻 [Code](https://github.com/OSU-NLP-Group/ARLoop)

ARLoop is a fully-open pretrained looped transformer with attention residuals. This release
contains the checkpoints of the 1.2B, 500B-token training run.
It is a **base pretrained model**, intended for text completion, which can be used for mid-training and post-training.

![Figure 1: ARLoop benchmark results](assets/figure1.png)

## Load and generate

```bash
pip install 'transformers==5.12.1' safetensors accelerate
# Install the PyTorch build appropriate for your CUDA environment separately.
```

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "osunlp/ARLoop-1.2B"  # Or the local ARLOOP_hf directory.
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id, trust_remote_code=True, dtype=torch.bfloat16,
).to("cuda").eval()
inputs = tokenizer("The capital of France is", return_tensors="pt").to("cuda")
with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))
```

Weights load in BF16 for inference. The stored checkpoint retains the original FP32 weights; routing and normalization statistics use FP32 internally.

## Architecture

| Property | Value |
| --- | --- |
| Physical transformer layers | 24, shared across 3 recurrent passes |
| Effective transformer depth | 72 |
| Hidden / MLP intermediate width | 2048 / 5632 |
| Attention / KV heads | 16 / 4 |
| Context length | 4096 |
| RoPE base | 50000 |
| Attention-residual block size | 24 sublayers |
| Recent residual blocks | 3, with input embedding retained |
| Value embeddings | Shared across alternating physical layers |
| Tokenizer vocabulary | 49152, SmolLM2 tokenizer used in training |
| Output head | Tied to token embeddings |

## License

Model release: [Apache License 2.0](https://www.apache.org/licenses/LICENSE-2.0).

## Training checkpoints

`main` contains the final checkpoint. Each revision below is independently loadable
with `revision="step-010000"` (use the same revision for the model and tokenizer), 
where 10000 denotes the number of optimization steps.

| Optimizer step | Training tokens | Revision |
| ---: | ---: | --- |
| 10,000 | 20,971,520,000 | [step-010000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-010000) |
| 20,000 | 41,943,040,000 | [step-020000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-020000) |
| 30,000 | 62,914,560,000 | [step-030000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-030000) |
| 40,000 | 83,886,080,000 | [step-040000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-040000) |
| 50,000 | 104,857,600,000 | [step-050000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-050000) |
| 60,000 | 125,829,120,000 | [step-060000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-060000) |
| 70,000 | 146,800,640,000 | [step-070000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-070000) |
| 80,000 | 167,772,160,000 | [step-080000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-080000) |
| 90,000 | 188,743,680,000 | [step-090000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-090000) |
| 100,000 | 209,715,200,000 | [step-100000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-100000) |
| 110,000 | 230,686,720,000 | [step-110000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-110000) |
| 120,000 | 251,658,240,000 | [step-120000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-120000) |
| 130,000 | 272,629,760,000 | [step-130000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-130000) |
| 140,000 | 293,601,280,000 | [step-140000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-140000) |
| 150,000 | 314,572,800,000 | [step-150000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-150000) |
| 160,000 | 335,544,320,000 | [step-160000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-160000) |
| 170,000 | 356,515,840,000 | [step-170000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-170000) |
| 180,000 | 377,487,360,000 | [step-180000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-180000) |
| 190,000 | 398,458,880,000 | [step-190000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-190000) |
| 200,000 | 419,430,400,000 | [step-200000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-200000) |
| 210,000 | 440,401,920,000 | [step-210000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-210000) |
| 220,000 | 461,373,440,000 | [step-220000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-220000) |
| 230,000 | 482,344,960,000 | [step-230000](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-230000) |
| 238,417 (final) | 499,996,688,384 | [main](https://huggingface.co/osunlp/ARLoop-1.2B/tree/main) / [step-238417](https://huggingface.co/osunlp/ARLoop-1.2B/tree/step-238417) |

All 23 scheduled intermediate checkpoints from 10k through 230k are available.

## Citation

```bibtex
@misc{yao2026arloop,
  title={{ARLoop}: Fully-Open Pre-trained Looped Transformers with Attention Residuals},
  author={Yuekun Yao and Yupei Du and Dawei Zhu and Huan Sun},
  year={2026},
  eprint={2611.05234},
  archivePrefix={arXiv},
  primaryClass={cs.CL},
  url={https://arxiv.org/abs/2611.05234}
}
```