DESIGN.md

maxact-fast

Train a model to generate text that maximally activates an arbitrary linear direction in its own residual stream, conditioned activation-oracle style (direction injected at layer 1, maximized at layer L). The bet: if we pretrain on a huge, diverse set of directions — not just SAE features — the model learns a general residual-direction→text inverter that transfers to SAE features for free (cross-uplift).

The direction firehose (why this beats SAE-only conditioning)

SAE conditioning caps you at ~65k features and inherits the SAE's blind spots. Instead we mint an effectively unlimited, diverse supply of concept directions straight from a corpus:

1. Embed a large corpus (UltraFineWeb, streamed) with a fast embedding model (BGE) — cheap.
2. Cluster the embeddings into K≈10k clusters (faiss k-means, GPU). Each cluster ≈ a concept.
3. For a training example, sample cluster A vs cluster B, train a quick linear probe on the policy model's layer-L residuals to separate A from B. The probe weight vector w is the conditioning direction — same space as an SAE encoder column, so it drops straight into the existing inject@1 / read@L machinery.
4. SFT target = the few real corpus texts closest to cluster A's centroid. "Given direction w, write text like this."
5. Pretrain on millions of these (shuffled). Then RL (Dr. GRPO, no KL) on cluster-probes and/or SAE features, and measure whether cluster-pretraining lifts SAE-feature performance.

Probe residuals are cached once per cluster (the only expensive Qwen3 pass), so step 3-4 mint millions of (direction, targets) examples cheaply.

Objective & injection (unchanged from the activation-oracle line)

• Inject unit(w) * ‖resid‖ * coeff into the layer-1 residual at a ? marker (norm-matched add).
• Reward = max over token positions of (x_L − μ)·unit(w), measured on the STANDALONE re-tokenized generation through the clean base model (no adapter, no injection). Same fair protocol as before.

RL = Dr. GRPO, done right, no KL

• advantage = r − group_mean — no /std (the removed Dr. GRPO difficulty bias).
• loss = Σ_tokens(−min(ratio·A, clip(ratio)·A) · mask) / total_tokens_in_batch (global constant normalizer, no per-response length bias).
• No KL term. Stability comes from lr + grad-clip + fluency/distinct gates (optional), not a leash.
• Rollouts via vllm-lens (SteeringVector(norm_match=True) == our injection), ~5× vs HF generate.
• Reuse vLLM generation logprobs as old_logps (one on-policy update/batch → ratio≈1, no extra fwd).

Pipeline stages (each a CLI in scripts/)

1. embed_cluster — stream corpus → BGE embeddings → faiss k-means K clusters → cluster assignments.
2. cache_resids — per cluster, run M member texts through Qwen3, cache layer-L mean residuals.
3. build_data — sample cluster pairs, train probes on cached resids, pick centroid targets → (direction, target_text) shards. Millions.
4. pretrain — SFT the generator (AO/LoRA) on the direction→text firehose. Massive.
5. finetune — final polish on per-feature MAX-activating CORPUS examples (cleanest targets).
6. rl — Dr. GRPO (no KL), vllm-lens rollouts, on cluster-probes and/or SAE features.
7. eval — beat-corpus / normalized-activation on held-out SAE features (cross-uplift test).
8. interface — tiny gradio app: type/pick a direction, watch generations + live activation.

Speed principles

bf16 + flash-attn everywhere; faiss-GPU k-means; HF streaming (never materialize the corpus); cache the one expensive Qwen3 residual pass; vLLM-lens rollouts; batchmax forwards on H200; torch.compile the training step. Target: pretraining data mine + pretrain in a day, not a week.