DJLougen/nanbeige-nextlat-experiment

Nanbeige NextLat Experiment

Two experimental continuations of Nanbeige/Nanbeige4.2-3B-Base: a plain control backbone and a "loop-time" backbone, plus loop-prediction heads for each.

What it is: a research artifact studying whether the second pass ("loop 2") of the model's repeated-layer execution can be predicted from the first and skipped for faster inference.

What's in this repo

Path	Contents
backbones/control/	control backbone (plain continued pretraining, ~100M tokens)
backbones/looptime/	loop-time backbone (continued pretraining + loop-prediction objective)
heads/	inference-only loop-transition heads (one per backbone)
tokenizer/	tokenizer + original config/code from upstream
code/	training + evaluation code (MIT)
configs/	experiment configs
metrics/	measured results (JSON)
provenance/	checksums, training metadata, data lineage

Result in one line: loop-2 becomes substantially more predictable after loop-time training, but skipping it still degrades generation. No speedup is demonstrated.

Quick start

Requires transformers>=4.44,<5 (tested with 4.57.6), PyTorch 2.x, sentencepiece, protobuf, and huggingface_hub.

import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer

# trust_remote_code does not resolve inside subfolders, so download to a local dir first
local = snapshot_download("DJI" "DJLougen/nanbeige-nextlat-experiment", allow_patterns=["backbones/looptime/*", "tokenizer/*"])
# local dir contains backbones/looptime/ and tokenizer/

model = AutoModelForCausalLM.from_pretrained(
    f"{local}/backbones/looptime",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,               # required: custom Nanbeige modeling code
).cuda()

tokenizer = AutoTokenizer.from_pretrained(f"{local}/tokenizer")

ids = tokenizer("The capital of France is", return_tensors="pt").input_ids.cuda()
out = model.generate(ids, max_new_tokens=32, do_sample=False)
print(tokenizer.decode(out[0]))

Loading a loop-prediction head

The heads predict loop-2 hidden states from loop-1 states. Architecture: LayerNorm(3072) → Linear(3072, 3072) → GELU → Linear(3072, 3072).

import torch, sys
sys.path.append("path/to/code")  # from this repo's code/ directory
from eval.train_head_streaming import make_head  # LayerNorm-Linear-GELU-Linear stack

ckpt = torch.load("heads/looptime_head_inference.pt", map_location="cpu", weights_only=True)
head = make_head(hidden_size=3072, device="cpu")
head.load_state_dict(ckpt["head_state"])
head.eval()

Training

Both backbones share the same data (FineWeb-Edu sample-10BT), sequence length 4096, and roughly 100M-token budget; the looptime arm adds an auxiliary loss that penalizes a linear map's failure to predict loop-2 hidden states from loop-1.

Arm	Tokens seen	Objective
control	100,004,702 (step 12,213)	plain LM loss
looptime	99,942,400 (step 12,200)	LM loss + loop-prediction aux (KL 0.15, warmup 5%)

The control head was retrained to match: 6,000 AdamW steps, batch 16,384, float32, frozen backbone, teacher-argmax CE + 0.3·MSE, lr 3e-4.

Results

Paired evaluation on 1,024 held-out positions (same inputs for both arms, seed 1729):

Metric	Control	Loop-time
Loop-2 prediction agreement (argmax)	57.7%	71.1%
Hidden-state cosine	0.895	0.911
Backbone CE (nats)	2.070	2.197

Skipping loop-2 does not work. Skipping runs ~1.96× faster but free-running generation agrees with full execution on only 3–11% of tokens, and outputs degrade to repetition on both backbones. Loop-time training improved predictability but not enough for naive skipping. See metrics/ for full protocols.

Limitations

• Small-scale continued pretraining (~100M tokens); these are research artifacts, not improved base models.
• Loop-time arm trades a small amount of CE quality (+0.13 nats) for predictability.
• Skip timing is batch-1 greedy with no KV cache and full-prefix recomputation; it is not production throughput or speculative-decoding acceptance.
• Hidden-state caches use a positional split; document-level independence was not established.
• The looptime head predates this comparison and lacks saved optimizer/RNG state, so its training provenance is less complete.

Licenses

• Model weights and upstream modeling/tokenizer code: Apache-2.0 (see LICENSE, UPSTREAM_NOTICE.md)
• Experiment code: MIT (see code/LICENSE)