PonderLM-2

TL;DR. Chain-of-Thought scales test-time compute by generating extra
tokens. PonderLM-2 does the same at pretraining time, but in
continuous space: before predicting each next token the model first
emits a few latent thoughts — extra last-hidden-state vectors — and
feeds them back into itself. Result: PonderLM-2-Pythia-1.4B (300 B Pile
tokens) beats vanilla Pythia-2.8B at equal inference flops, on language
modeling and a range of downstream tasks. More latent thoughts, lower
loss.

   vanilla:      x₁ ──► x₂ ──► x₃ ──► x₄

   PonderLM-2:   x₁ ──► z₁ ──► x₂ ──► z₂ ──► x₃ ──► z₃ ──► x₄ ──► z₄
                       z_i = latent thought emitted before predicting x_{i+1}

What's in this repo

A slim fork of LLaMA-Factory
covering Llama and GPT-NeoX. PonderLM-2, three baselines, and a vanilla
reference all share the same CLI and switch with one flag:

Released checkpoints on the Hugging Face Hub:

• 🤗 zeng123/PonderLM-2-Pythia-1.4b
• 🤗 zeng123/PonderLM-2-Pythia-410m

Repository tour

PonderLM-2/
├── scripts/                      # one script per row of Table 3 + quickstart
├── llama_config/                 # Llama configs (125m, 410m, 834m, 1.4b, 2.8b)
├── pythia_config/                # Pythia configs (410m)
├── examples/deepspeed/           # ZeRO configs used in the paper
├── data/dataset_info.json        # minipile, smallpile, testpile, uint16smallpile
└── src/llamafactory/model/
    ├── llama_patch.py                                   # method swap entry point
    └── modeling/
        ├── modeling_llama.py                            # PonderLM-2 (Llama, w/ KV-cache)
        ├── modeling_llama_orin.py                       # PonderLM   (Llama)
        ├── modeling_llama_loop.py                       # Loop       (Llama)
        ├── modeling_llama_pause.py                      # Pause      (Llama)
        ├── modeling_gpt_neox.py                         # PonderLM-2 (GPT-NeoX)
        ├── modeling_gpt_neox_addhidden.py               # PonderLM   (GPT-NeoX)
        ├── modeling_gpt_neox_addhidden_weightshare.py   # Loop       (GPT-NeoX)
        └── modeling_gpt_neox_addpausetoken.py           # Pause      (GPT-NeoX)

Install

git clone https://github.com/hzby/PonderLM-2.git
cd PonderLM-2
conda create -n ponderlm2 python=3.10 -y
conda activate ponderlm2
pip install -e ".[torch,metrics,deepspeed]"
pip install wandb
conda install cuda=12.1.0 -c nvidia -y
pip install flash-attn==2.7.2.post1 --no-build-isolation

Tested on PyTorch 2.4 + CUDA 12.1, 8×A100 / 8×H100 / 8×RTX4090.

Quickstart — Pythia-70m on minipile

A small pretraining run on 8 GPUs to confirm the install works before
committing to the larger Table 3 runs.

bash scripts/quickstart_ponderlm2_pythia_70m.sh

Behind the scenes:

FORCE_TORCHRUN=1 llamafactory-cli train \
    --model_name_or_path EleutherAI/pythia-70m \
    --stage pt --do_train --finetuning_type full --train_from_scratch true \
    --dataset minipile --template default --cutoff_len 2048 \
    --num_train_epochs 1.0 \
    --per_device_train_batch_size 8 --gradient_accumulation_steps 2 \
    --learning_rate 1.0e-3 --weight_decay 0.01 --warmup_ratio 0.01 \
    --lr_scheduler_type cosine_with_min_lr \
    --lr_scheduler_kwargs '{"min_lr_rate":0.1}' \
    --bf16 --flash_attn fa2 --disable_gradient_checkpointing true \
    --deepspeed examples/deepspeed/ds_z0_config.json \
    --patch_method ours --recurrent_model true --scale_embeds true \
    --interpolation false \
    --num_latent_thoughts 1 \
    --num_jacobi_iterations 3 \
    --random_jacobi_iterations true \
    --output_dir saves/minipile/ponderlm2_pythia_70m

For a vanilla baseline next to it, replace --patch_method ours with
--patch_method vanilla and drop --recurrent_model true; the rest of the
ponder flags become no-ops. After a few thousand steps the PonderLM-2
training loss settles below the vanilla curve and stays there for the rest
of the run.

Smaller-compute companion (Llama-410m on smallpile)

A lighter, cheaper version of the Table 3 experiments below — same
recipe, just lr=3e-4 and the 410m base. Useful for sanity-checking the
method or for users without enough compute to run the full 1.4b sweep.

The corpus is the same smallpile (uint16-tokenised) used in Table 3:

mkdir -p data/tokenized_data
huggingface-cli download hyq718/uint16smallpile \
    --repo-type dataset --local-dir data/tokenized_data/smallpile

Three runs to compare PonderLM-2 against Loop=3 (extra passes through
the layer stack) and a 2× parameters vanilla baseline (the 410m
architecture with the layer count doubled from 24 to 48, ≈ 820M params,
config at llama_config/410m_x2depth):

bash scripts/train_ponderlm2_llama_410m.sh           # ours, 410m
bash scripts/train_loop_llama_410m.sh                # Loop, K=3, 410m
bash scripts/train_vanilla_llama_410m_x2depth.sh     # vanilla, 410m × 2 depth (≈820M)

We observe that the 410m PonderLM-2 run ends up with lower loss than
both the K=3 Loop baseline and the doubled-depth 820M vanilla model at
the same step count.

Reproducing Table 3 (Llama-1.4b on smallpile)

Table 3 of the paper compares PonderLM-2 against Pause Token, Loop, and
PonderLM, all 1.4 B-parameter Llamas trained on smallpile for 50 000
steps at global batch 256, seq-len 2048 (≈ 26 B tokens) on an 8-GPU node.
The repo ships one script per row, plus a vanilla reference.

The corpus is pre-tokenised as uint16 and lives at
hyq718/uint16smallpile:

mkdir -p data/tokenized_data
huggingface-cli download hyq718/uint16smallpile \
    --repo-type dataset --local-dir data/tokenized_data/smallpile

Then launch one row at a time:

bash scripts/train_ponderlm2_llama_1_4b.sh   # ours
bash scripts/train_ponderlm_llama_1_4b.sh    # PonderLM
bash scripts/train_loop_llama_1_4b.sh        # Loop
bash scripts/train_pause_llama_1_4b.sh       # Pause Token
bash scripts/train_vanilla_llama_1_4b.sh     # vanilla reference

Each script picks up three environment variables (MODEL_CFG, TOKENIZED,
OUTPUT_DIR) so you can override the model config, the tokenised-data
path, or where checkpoints get written without touching the script.
Switch to GPT-NeoX with MODEL_CFG=pythia_config/1.4b (or 410m).
Multi-node works through the LLaMA-Factory launcher — just prefix the
command with torchrun --nnodes ... or your scheduler's launcher.

Hyperparameters

Everything else (optimiser, scheduler, DeepSpeed, FlashAttention, data flags)
is plain LLaMA-Factory; see the scripts for the values used in the paper.

Inference

The released GPT-NeoX checkpoints are wired up via auto_map, so a normal
from_pretrained does the right thing:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

ckpt = "zeng123/PonderLM-2-Pythia-1.4b"
tok = AutoTokenizer.from_pretrained(ckpt)
model = AutoModelForCausalLM.from_pretrained(
    ckpt, torch_dtype=torch.bfloat16, trust_remote_code=True,
).cuda()

prompt = "The mitochondria is "
out = model.generate(
    **tok(prompt, return_tensors="pt").to(model.device),
    max_new_tokens=64, use_cache=True,
)
print(tok.decode(out[0], skip_special_tokens=True))

For Llama checkpoints trained from this repo, patch the class first:

from llamafactory.model.llama_patch import patch_llama_ponderlm2
patch_llama_ponderlm2()
# from_pretrained as usual

The Llama path supports the standard KV-cache, so model.generate works
the same way as for any HuggingFace transformer.

Citation

@article{zeng2025ponderlm,
  title={Ponderlm-2: Pretraining llm with latent thoughts in continuous space},
  author={Zeng, Boyi and Li, He and Song, Shixiang and Wang, Yixuan and Wang, Zitong and He, Ziwei and Wang, Xinbing and Lin, Zhouhan},
  journal={arXiv preprint arXiv:2509.23184},
  year={2025}
}

Acknowledgements

Built on LLaMA-Factory
(Apache-2.0). The PonderLM baseline is adapted from
LUMIA-Group/PonderingLM.

Contact

Open a GitHub issue, or email
[email protected].