99kj/where-winter-goes

Where Winter Goes: GRAM checkpoints

The 45 models from Where Winter Goes. It is a study of whether steering vectors can bring back a concept that GRAM (gradient routing with auxiliary modules) removed. The concept is winter, in a synthetic corpus of templated children's stories where one training window in six is a winter story. GRAM sends labelled winter text through a small module in each layer and deletes the module before the model ships. p_as is the chance that a batch's labelled winter text also trains the rest of the model.

These are research artifacts for studying capability removal and steering. They are not general-purpose language models: the vocabulary has 123 words and the context is 48 tokens.

What's here

Models	pas	Seeds	Files
GRAM	0, 0.1, 0.25, 0.75, 0.9, 0.95, 0.98	0–2	models/t7_winter_gram_pas<p>_seed<s>
GRAM	0.5 (the default)	0–5	models/t7_winter_gram_seed<s>
GRAM	1	0–5	models/t7_winter_gram_pas1_seed<s>
Dense, trained on everything		0–5	models/t7_winter_dense_seed<s>
Filtered, winter windows removed		0–5	models/t7_winter_filter_seed<s>

Each model has <tag>.safetensors, a state dict for track7_concepts.LM, and <tag>.json, its training config and the metrics logged at the end of training. vocab.json lists the 123 tokens in index order, and config.json describes the architecture.

The GRAM checkpoints still contain the module. Run them with mask=(0,) for the shipped model (module deleted) and mask=(1,) for the full model with the module attached. The dense and filtered models share the architecture, but their module was never trained and stays at its random initialisation; always run them with mask=(0,).

Architecture and training

A 4-layer pre-norm transformer: width 160, 4 attention heads, learned positional embeddings, a 48-token context, and a final layer norm and linear output layer. Each feed-forward block is a core MLP (160 → 640 → 160, GELU) plus a GRAM module (160 → 80 → 160, GELU), about 1.39M parameters in all, 103,360 of them in the modules. Training ran for 4 epochs over 30,000 ordinary and 6,000 winter windows, with batch size 64, AdamW (lr 3e-4, weight decay 0.01), and one optimizer per parameter group (core, module).

GRAM routing, per batch: winter windows run with the module on, always train the module, and train the core with probability p_as. Ordinary windows always train the core, and in half the batches run with the module on and train it too. The filtered model never sees a winter window. The dense model sees everything.

Load a model

The model class lives in the GitHub repository:

git clone https://github.com/kilojoules/where-winter-goes && cd where-winter-goes
pip install -r requirements.txt

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
import track7_concepts as T

path = hf_hub_download("99kj/where-winter-goes", "models/t7_winter_gram_pas1_seed0.safetensors")
m = T.LM()
m.load_state_dict(load_file(path))
m.eval()
x = torch.tensor([[T.TOK[w] for w in ["one", "day", "in"]]])
with torch.no_grad():
    for mask in [(0,), (1,)]:  # module deleted, module attached
        p = torch.softmax(m(x, mask=mask)[:, -1], -1)[0, T.TOK["winter"]]
        print(mask, float(p))  # (0,) 0.0043, (1,) 0.36

python fetch_models.py in the repository downloads every checkpoint into results/*.pt, the format its evaluation scripts read.

What the models do

Unsteered. The two "Starts winter" columns are ranges over seeds; loss and facts are means over seeds. "Starts winter" is the probability that the next word after "one day in" is "winter". "Facts" is how often the model completes a winter story with the right item for its activity (16 activities, 4 items; a constant guess scores 0.25).

Model	Starts winter, module deleted	Starts winter, module attached	Loss on winter stories (nats/token)	Facts
GRAM pas 0	3.1e-7 – 6.9e-7	1.3e-4 – 0.0016	4.74	0.00
GRAM pas 0.1	6.3e-7 – 2.3e-6	0.0013 – 0.23	4.29	0.07
GRAM pas 0.25	9.4e-7 – 5.9e-6	0.0073 – 0.16	3.80	0.17
GRAM pas 0.5	6.8e-7 – 1.9e-6	0.026 – 0.21	3.06	0.63
GRAM pas 0.75	2.3e-6 – 6.8e-6	0.12 – 0.41	1.50	1.00
GRAM pas 0.9	2.4e-5 – 1.7e-4	0.15 – 0.36	0.80	1.00
GRAM pas 0.95	2.9e-4 – 0.0011	0.17 – 0.37	0.57	1.00
GRAM pas 0.98	0.0016 – 0.0030	0.17 – 0.37	0.48	1.00
GRAM pas 1	0.0029 – 0.015	0.13 – 0.55	0.45	1.00
Dense	0.093 – 0.19		0.35	1.00
Filtered	7.9e-6 – 3.6e-5		3.50	0.00

At high p_as the shipped GRAM core knows winter stories but rarely starts one. A push toward the word "winter" makes it start one, but the stories that follow are mostly not winter stories. At p_as 0.5 no steering vector brings winter back. The summary page and the repository have the full results, the pre-registered tests and the limits.

License

MIT, for the weights and the code.