ceselder/maxact-fast-pretrain-tokmax-100k

maxact-fast — pretrain (SL) checkpoint, k=100k token-max run

A Qwen3-8B LoRA adapter that inverts a residual-stream direction → text that maximally activates it. You inject a direction v (any layer-27 residual direction — an SAE decoder/encoder column, a linear probe, a mean-difference vector, …) at a marker token, and the model generates text whose layer-27 activations line up with v.

This is the supervised-pretrain (SL) checkpoint (pre-RL) from the k=100k token-max run — the one used for the SL-vs-RL scaling analysis. Good for qualitative exploration; a separate RL checkpoint pushes zero-shot SAE activation ~+80% higher.

How it was trained

• Base: Qwen/Qwen3-8B. Adapter: LoRA r=64, α=16, rsLoRA, all linear layers.
• Data (1.5M pairs): k-means (k=100k, whiten+L2-norm space) over mean-pooled layer-27 residuals of ~113M corpus spans → 300k linear probes (auc≈1.0). For each probe, a token-max retrieval over 30M spans finds the span containing its single highest-activating token → that span is the target text (5 targets/probe).
• Objective: plain supervised CE — given v injected at the marker, produce the target text.
• Optim: AdamW lr 3e-5, OneCycle 2% warmup, 1 epoch (1465 steps), global batch 1024 (128/GPU × 8×B300), max_seq 192. Loss 4.3 → ~2.6.

How injection works (this is the whole interface)

• Prompt: the instruction "Please produce a string of text that triggers the following direction maximally:" rendered with the chat template (add_generation_prompt=True), then a single ? marker appended right after. The marker is the last prompt token (index 26 of the 27-token prompt).

• Inject at decoder layer 1 (0-indexed), norm-matched additive, only at the marker position:

  h[marker] += unit(v) · ‖h[marker]‖ · coeff        # coeff = 1.0
  

  i.e. add the unit direction scaled to the residual's own norm. v is normalized internally; sign matters. (During KV-cache decode steps the hook is a no-op — the marker was injected at prefill.)

• Score / read at layer 27: the quantity being maximized is max_t (h27[t] · unit(v)) over the generated tokens (re-forward the generated text standalone through the base model).

• Sampling best-of-N (temperature 1.0) noticeably beats greedy.

Quick start

demo.py (included) uses the two helper files bundled here (prompts.py, inject.py — the exact training-time code). Minimal version:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from prompts import build_prompt_ids            # bundled in this repo
from inject import make_inject_hook, get_layer  # bundled in this repo

REPO, MODEL, INJECT_LAYER, COEFF = "ceselder/maxact-fast-pretrain-tokmax-100k", "Qwen/Qwen3-8B", 1, 1.0
tok  = AutoTokenizer.from_pretrained(MODEL)
base = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="cuda")
model = PeftModel.from_pretrained(base, REPO).eval()

v = torch.randn(4096)                            # <-- your layer-27 direction (sign matters)
prompt_ids, mpos = build_prompt_ids(tok)         # mpos = [marker index]
ids = torch.tensor([prompt_ids], device="cuda")
hook = make_inject_hook([v.reshape(1, -1).cuda()], [mpos], COEFF, "cuda", torch.bfloat16, mode="add")
h = get_layer(model, INJECT_LAYER).register_forward_hook(hook)
try:
    out = model.generate(ids, max_new_tokens=64, do_sample=True, temperature=1.0, top_p=0.95)
finally:
    h.remove()
print(tok.decode(out[0, len(prompt_ids):], skip_special_tokens=True))

Requires transformers>=4.44, peft, torch. Constants: INJECT_LAYER=1, read/score layer 27, D_MODEL=4096, MARKER=" ?", STEER_COEFF=1.0.