LatentLoop

Pretrained checkpoints for LatentLoop: Scaling Looped MoE Models with Asymmetric Compute Allocation.

Paper · Code · Project page

LatentLoop compresses the sequence during the first visit through the repeated middle block, retains full sequence resolution in the second visit, and reallocates saved computation to increase model width. MoE blocks are shared across visits, while attention blocks have separate parameters. A learned scaled decoder calibrates the latent increment before its causally aligned writeback.

This repository contains the largest-budget LatentLoop endpoint at each of the three model scales. Each subfolder is a separate model checkpoint; the files are not shards of one combined model.

Checkpoint	Reference token budget	Actual prediction tokens	Validation NLL	0-shot Avg (%)	5-shot Avg (%)
1.3B	40B	40.0975B	2.1393	50.48	51.91
3B	70B	69.6538B	1.9771	56.69	58.51
7B	140B	139.7123B	1.8482	61.57	63.40

Size labels denote approximate total parameter counts. Reference token budgets identify the matched-FLOPs settings; actual prediction-token counts differ slightly because training steps are adjusted to match computation. Validation NLL uses the fixed 100,003,840-token validation panel. Avg is the unweighted mean of nine benchmarks: LAMBADA OpenAI, LAMBADA Standard, HellaSwag, PIQA, WinoGrande, ARC-Easy, ARC-Challenge, SciQ, and MMLU continuation. Both shot settings use no prefix token, few-shot seed 1234, and one evaluation per setting. Detailed scores and the evaluation protocol are in evaluation_results.json.

Load and generate

Use the native PyTorch implementation included here or in the code repository. Install PyTorch 2.9.0 and the optional inference dependencies:

git clone https://github.com/lcqysl/LatentLoop.git
cd LatentLoop
pip install -r requirements.txt -r requirements-inference.txt
python inference.py --model yhx12/LatentLoop --subfolder 1.3B-40B \
  --prompt "The capital of France is" --max-new-tokens 32

Select 3B-70B or 7B-140B with --subfolder to load another scale. The loader downloads only the selected subfolder. It validates the complete tensor index and loads FP32 parameters, using BF16 autocast for CUDA inference. The default uses PyTorch SDPA, grouped MoE matrix multiplication, and a PyTorch SwiGLU implementation, so FlashAttention is optional. The paper's FlashAttention/SwiGLU kernels can be selected with --attention-backend flash --mlp-backend flash after installing compatible extensions. Kernel choices can affect throughput and numerical rounding.

For Python use:

import torch
from inference import load_pretrained, generate

model, tokenizer = load_pretrained("yhx12/LatentLoop", subfolder="1.3B-40B", device="cuda")
tokens = torch.tensor([tokenizer.encode("The capital of France is")], device="cuda")
with torch.autocast("cuda", dtype=torch.bfloat16):
    output = generate(model, tokens, max_new_tokens=32, eos_token_id=tokenizer.eos_id())
print(tokenizer.decode(output[0].tolist()))

The included loader manages the separate KV state for each visit and the pending pair/writeback state needed for incremental LatentLoop decoding. The context length is 4096 tokens. These are base pretrained language models for text continuation and architecture research. No instruction tuning or chat template is applied.

Training and files

All models were trained from scratch on a fixed FineWeb-Edu sample350BT stream, with exact token-document deduplication, a fixed shuffle, and whole-document separation of training and validation. They use the 32K TinyLlama SentencePiece tokenizer, no added BOS, and one EOS per training document. The tokenizer's automatic BOS insertion is disabled in this release to match the experimental protocol; the SentencePiece model is unchanged. Data preparation and the matched model configurations are provided in the code repository.

Each checkpoint directory contains:

• model-*-of-*.safetensors and model.safetensors.index.json: model parameters in their original FP32 precision.
• model_config.json: the configuration for the native model.Config class.
• tokenizer/: the tokenizer and special-token configuration.
• training_summary.json: actual token counts, optimizer steps, validation NLL, and original checkpoint identity.
• checksums.json: SHA256 hashes of the released checkpoint files.

Optimizer state, random-number-generator state, and private training paths are not included. The released tensors are checked for finiteness and exact equality to the source checkpoint. Loading and forward computation are checked against the frozen implementation used to train each endpoint.

License

LatentLoop weights and original code are released under the Apache License 2.0. Third-party source portions retain their notices and licenses; see NOTICE. The included tokenizer originates from TinyLlama.

Citation

@article{he2026latentloop,
  title={LatentLoop: Scaling Looped MoE Models with Asymmetric Compute Allocation},
  author={He, Zefeng and Chai, Yunkai and Chen, Xiaoyang and Zhu, Tong and Wu, Zongliang and Luo, Yun and Li, Yafu and Cheng, Yu},
  journal={arXiv preprint arXiv:2610.19975},
  year={2026}
}