name-sort-pretrain-4L

A small GPT-2 language model (4 layers, hidden size 256, 4 heads, MLP size 1024, context length 256) pretrained from scratch on the sentences behind a synthetic name-sorting task. It is the pretrained base that is later fine-tuned on the task itself, as part of a benchmark of ground-truth circuits for evaluating circuit-discovery methods.

Training data

Luoxiaoxi/synthetic-name-index-pretrain: real sentence templates with every person name replaced by a synthetic 3-character name written as three single-character words (e.g. j u j), packed into 256-token blocks of <sentence> [EOS]. Only templates from the fine-tuning train split are used.

The task answer and its [ANS] separator never appear in pretraining. The fine-tuning format is <sentence> [ANS] <answer> [EOS], where the answer (task index) is the first-appearance position (0-4) of the lexicographically smallest name. Because this checkpoint has never seen [ANS], its predictions after [ANS] are not meaningful until it is fine-tuned. Paired clean/counterfactual task prompts are in Luoxiaoxi/synthetic-name-index-counterfactual.

Training: 1,500 steps, batch size 64, AdamW (learning rate 5e-4, weight decay 0.01), cosine schedule with 50 warmup steps, dropout 0.1, seed 42. Final training loss 3.90.

Tokenizer

The tokenizer is a custom word-level tokenizer, WhitespaceDigitTokenizer, included in this repo as ws_tokenizer.py (it only depends on transformers):

• text is split on whitespace only, and any token containing digits is further split into single digits (1912 -> 1 9 1 2);
• the vocabulary (37,139 entries) is every word of the training corpora plus the special tokens [UNK], [PAD], [EOS] and [ANS]; any other word maps to [UNK];
• decoding joins tokens with single spaces.

The corpus is pre-tokenized, so punctuation and clitics are separate words: write her , she and once ., not her, she and once. (those become [UNK]).

Loading

With trust_remote_code, AutoTokenizer loads the class from ws_tokenizer.py:

from transformers import AutoTokenizer, GPT2LMHeadModel

repo = 'Luoxiaoxi/name-sort-pretrain-4L'
tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = GPT2LMHeadModel.from_pretrained(repo)

Without remote code, download the file and use the class directly:

import importlib.util
from huggingface_hub import hf_hub_download
from transformers import GPT2LMHeadModel

repo = 'Luoxiaoxi/name-sort-pretrain-4L'
spec = importlib.util.spec_from_file_location('ws_tokenizer', hf_hub_download(repo, 'ws_tokenizer.py'))
ws_tokenizer = importlib.util.module_from_spec(spec)
spec.loader.exec_module(ws_tokenizer)
tokenizer = ws_tokenizer.WhitespaceDigitTokenizer.from_pretrained(repo)
model = GPT2LMHeadModel.from_pretrained(repo)

Example

import torch

text = 'When j u j does not believe her , she claims she heard z u z sing once .'
ids = tokenizer(text, add_special_tokens=False, return_tensors='pt').input_ids
print(tokenizer.convert_ids_to_tokens(ids[0]))  # one token per word; no [UNK]
with torch.no_grad():
    next_id = model(ids).logits[0, -1].argmax()
print(tokenizer.decode([int(next_id)]))

Files

• model.safetensors, config.json, generation_config.json: the final checkpoint.
• vocab.json, tokenizer_config.json, special_tokens_map.json, ws_tokenizer.py: the tokenizer.
• config.yaml, trainer_state.json, train_results.json, all_results.json, training_args.bin: the training run's config, log and results.