Coconut-0.2B-gsm8k-aug

🌐 Español · English

Our own Coconut (Hao et al., 2024) retrain on GSM8k-Aug, from GPT-2. Coconut reasons in a continuous latent space: a fixed number of hidden states are fed back as input embeddings before the model switches to text and decodes the answer.

Why this repo exists. There is no official public Coconut GSM8k checkpoint. We trained one so the latent-halting (PonderNet) work — UdeSA/ALST-0.2B-gsm8k-aug — has a plain Coconut arm to compare against and to warm-start its Coconut + halting variant from. This is a reproduction, not a new method: run under upstream's recipe, it reproduces the Coconut paper's own GSM8k figure for GPT-2.

Results

GSM8K test, n = 1,319, greedy decoding. Epoch selected on the 500-row whynlp validation split; test reported once.

leaf	model	test accuracy	depth	note
coconut/	Coconut, fixed latent depth K = 6	34.57%	K = 6	reproduces the paper's 34.1%
cot/	stage-0 chain-of-thought SFT	43.14%	text	the warm-start the curriculum starts from

K = 6 is c_thought 2 × max_latent_stage 3 (3 latent groups of 2 continuous thoughts each), Coconut's trained fixed depth — it has no halting, so this single operating point is the whole model.

Layout

Two leaves, each a self-contained snapshot (weights + tokenizer + summary.json):

• coconut/ — the reported Coconut model. The state dict is the Coconut wrapper's (keys prefixed base_causallm., vocab 50260 = GPT-2 + 3 latent markers <|start-latent|>, <|end-latent|>, <|latent|>). Load it with the Coconut class from the code, not from_pretrained.
• cot/ — the stage-0 CoT model, a standard GPT2LMHeadModel (vocab 50257) that loads directly with from_pretrained.

Loading

The stage-0 CoT model loads directly:

from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import snapshot_download

path = snapshot_download("UdeSA/coconut-0.2B-gsm8k-aug", allow_patterns="cot/*")
model = AutoModelForCausalLM.from_pretrained(f"{path}/cot")
tok = AutoTokenizer.from_pretrained(f"{path}/cot")

The Coconut model needs the Coconut class (see src/coconut/ in the project repo), which wraps the backbone and runs the latent loop:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import snapshot_download
from src.coconut.model import Coconut

path = snapshot_download("UdeSA/coconut-0.2B-gsm8k-aug", allow_patterns="coconut/*")
tok = AutoTokenizer.from_pretrained(f"{path}/coconut")
latent, start, end = (tok.convert_tokens_to_ids(t)
                      for t in ("<|latent|>", "<|start-latent|>", "<|end-latent|>"))
backbone = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
backbone.resize_token_embeddings(len(tok))
model = Coconut(backbone, latent, start, end, tok.eos_token_id)
model.load_state_dict(torch.load(f"{path}/coconut/pytorch_model.bin", weights_only=True))

Recipe

From GPT-2 (openai-community/gpt2), upstream's hyperparameters unchanged: AdamW lr 1e-4, weight decay 0.01, batch 32, c_thought 2, epochs_per_stage 3, max_latent_stage 3, pad_latent_to_max True, uniform_prob 0.0, 25 epochs per stage, seed 0, bf16. Two sequential runs: stage-0 CoT SFT (--cot True), then the latent curriculum warm-started from it (--coconut True). Data: GSM8k-Aug (385,620 train / 500 whynlp validation / 1,319 test).