ContextLab / ContextLab/llm-geometry

Geometry Lab: fine-tuning ignores a scratch-trained model's own vocabulary

Open
#6 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
TypeScript
Stars
1
Forks
0
PR merge metrics
No merged PRs in 30d

Description

Found by an adversarial review of the feature-005 documentation: a panel note claimed
fine-tuning "keeps the existing vocabulary". It does not, and the underlying behaviour is
wrong rather than the note being merely sloppy.

## What happens

`train from scratch` mints a model with a **fresh 1000-word vocabulary** built from the
user's text, so its token ids mean different words than the shipped checkpoint's.
Fine-tuning that model then encodes the fine-tuning text with the **shipped** vocabulary:

- `code/backend/src/llm_geometry/geo/finetune.py:157` — `ids = get_tokenizer().encode_stream(text)`
(the module-level canonical tokenizer, never `tokenizer_for(base)`)
- `code/frontend/src/lib/staticClient/geo.ts` — same, and its own docstring says
"Fine-tuning keeps the shipped vocabulary"

So gradients land on rows whose ids mean something else entirely.

It also **loses** the base's vocabulary on the way out:
`finetune.py:207` calls `save_weight_set(new_ws, source="finetuned", store=store)` with no
`vocab_json`, so `tokenizer_for()` falls back to the canonical tokenizer for the result —
and every dot on the sphere is relabelled with the shipped words.

## Repro

1. Train from scratch on any text with ≥1000 distinct types.
2. Fine-tune the resulting model on a paragraph.
3. Hover the dots: they are labelled with the shipped checkpoint's words, not yours.

## Fix sketch

- `finetune.py`: tokenize with `tokenizer_for(base_token, store)`, and pass the base's
`vocab_json` through to `save_weight_set` so the result keeps its own words.
- Mirror in `staticClient/geo.ts` via the engine's `tokenizerFor(base)`.
- Test: scratch-train, fine-tune, assert the fine-tuned model's vocabulary equals the
scratch model's and that a known word round-trips to the same id.

Documented in the meantime: the Info tab and the fine-tune panel both now say outright that
fine-tuning uses the shipped vocabulary, and link here.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with code/backend/src/llm_geometry/geo/finetune.py at lines 157 and 207, then compare the matching flow in code/frontend/src/lib/staticClient/geo.ts. Run the scratch-training and fine-tuning reproduction described in the issue. Done means fine-tuning preserves the scratch model's vocabulary and a known word round-trips to the same token id in both paths.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, typescript
Domain
backend, frontend, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
70/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.