PonderLM-2-Pythia-1.4b

Pythia-1.4b architecture pretrained with PonderLM-2, the method introduced in PonderLM-2: Pretraining LLM with Latent Thoughts in Continuous Space (ICML 2026 Spotlight).

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: this 1.4B model trained on 300 B Pile tokens beats vanilla Pythia-2.8B at equal inference flops, on language modelling and a range of downstream tasks.

   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}

• Code: LUMIA-Group/PonderLM-2
• Paper: arXiv:2509.23184
• Sibling checkpoint: zeng123/PonderLM-2-Pythia-410m

---

Usage

The model ships with a custom modeling_gpt_neox.py that runs the pondering forward pass. Loading via AutoModelForCausalLM requires trust_remote_code=True:

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))

---

Model details

---

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 top of the Pythia training stack and LLaMA-Factory. The PonderLM baseline implementation is adapted from LUMIA-Group/PonderingLM.