Phase 1 release: InterpGPT matched-pair checkpoint

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,87 @@
+---
+license: mit
+library_name: transformers
+tags:
+  - interpretability
+  - mechanistic-interpretability
+  - task-decomposition
+  - small-language-model
+  - transformer-lens
+pipeline_tag: text-generation
+---
+
+# InterpGPT — ADHD Model (23M)
+
+Part of the **InterpGPT** matched-pair release. This is the **ADHD** model;
+its counterpart is
+[`connaaa/interpgpt-standard-23M`](https://huggingface.co/connaaa/interpgpt-standard-23M).
+Both models share identical architecture and training recipe; only the
+training data distribution differs.
+
+**ADHD variant training data**: task decompositions broken into smaller steps
+with interleaved micro-regulation actions ("sip water", "deep breath",
+"close eyes briefly", "quick stretch", "pause").
+
+| | Value |
+|---|---|
+| Parameters | 23,471,104 |
+| Layers | 6 |
+| Heads | 8 |
+| d_model | 512 |
+| d_head | 64 |
+| d_mlp (SwiGLU) | 1408 |
+| Vocab | 8192 (custom BPE) |
+| Context length | 512 |
+| Norm | RMSNorm (ε = 1e-6) |
+| Position | RoPE (half-half, base 10,000) |
+| Activation | SwiGLU |
+| Biases | none |
+| Tied input/output embeddings | yes |
+| Training tokens | ~25k steps on ADHD-variant task-decomposition corpus |
+
+## Headline findings (Phase 1)
+
+- **Structural head-position swap.** A step-layout-broadcast head lives at
+  **L3H0** in the standard model and at **L3H5** in the ADHD model.
+  Cross-model per-position attention profile cosine at the matched pair
+  **0.997**; same-index baseline **0.66** (0.663 for one pair; 0.643 for another).
+  Causal ablation confirms the functional identity: ablating L3H5 in the ADHD
+  model drops Spearman(task_complexity × step_count) from 0.83 → 0.78 (median
+  Δ = -0.055 across 5 seeds).
+- **Block-2 content circuit.** P(regulation token) at step-onset positions
+  jumps 17× between layer 1 and layer 2 (0.014 → 0.251). The standard model
+  never crosses 1% at any layer.
+- **High-specificity null-steering feature.** An ADHD-L2 SAE feature
+  (feat 2504) fires at 59% of ADHD step-onsets vs 0.03% of standard step-onsets
+  (~2000× cross-model asymmetry), yet **causal steering on its decoder
+  direction produces Δ within sampling noise under all four intervention
+  variants** (inject-std, subtract-adhd, zero-ablate, inject-upstream).
+  See the companion SAE repo
+  [`connaaa/interpgpt-sae-phase5`](https://huggingface.co/connaaa/interpgpt-sae-phase5).
+
+## Loading
+
+Identical to the standard variant. See
+[`connaaa/interpgpt-standard-23M`](https://huggingface.co/connaaa/interpgpt-standard-23M)
+for `AutoModel`, TransformerLens, and raw-TaskGPT examples, substituting the
+repo id.
+
+## Input format
+
+```
+<|task|>Clean the kitchen<|steps|>Step 1 text<|sep|>Step 2 text<|sep|>...<|end|>
+```
+
+## Reproduce the head-swap finding
+
+Open the Colab at `notebooks/InterpGPT_HeadSwap.ipynb`
+(https://github.com/cwklurks/interpgpt). Runs end-to-end on Colab free tier in
+under 15 minutes.
+
+## License
+
+MIT.
+
+## Citation
+
+See the standard model card.

config.json

@@ -0,0 +1,78 @@
+{
+  "vocab_size": 8197,
+  "max_seq_len": 512,
+  "n_layers": 6,
+  "n_heads": 8,
+  "d_model": 512,
+  "d_ff": 2048,
+  "dropout": 0.25,
+  "pad_id": 8196,
+  "bias": false,
+  "variant": "adhd",
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": null,
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "architectures": null,
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "task_specific_params": null,
+  "problem_type": null,
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": null,
+  "pad_token_id": 8196,
+  "eos_token_id": null,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "_name_or_path": "",
+  "transformers_version": "4.57.6",
+  "auto_map": {
+    "AutoConfig": "configuration_interpgpt.InterpGPTConfig",
+    "AutoModel": "modeling_interpgpt.InterpGPTModel"
+  },
+  "tf_legacy_loss": false,
+  "use_bfloat16": false,
+  "model_type": "interpgpt",
+  "output_attentions": false
+}

configuration_interpgpt.py

@@ -0,0 +1,37 @@
+"""
+HuggingFace PretrainedConfig for InterpGPT / TaskGPT.
+
+Mirrors gpt_model.GPTConfig but subclasses transformers.PretrainedConfig
+so `AutoConfig` / `AutoModel.from_pretrained(..., trust_remote_code=True)` work.
+"""
+from transformers import PretrainedConfig
+
+
+class InterpGPTConfig(PretrainedConfig):
+    model_type = "interpgpt"
+
+    def __init__(
+        self,
+        vocab_size: int = 8192,
+        max_seq_len: int = 512,
+        n_layers: int = 6,
+        n_heads: int = 8,
+        d_model: int = 512,
+        d_ff: int = 2048,
+        dropout: float = 0.1,
+        pad_id: int = 0,
+        bias: bool = False,
+        variant: str = "standard",
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.dropout = dropout
+        self.pad_id = pad_id
+        self.bias = bias
+        self.variant = variant
+        super().__init__(pad_token_id=pad_id, **kwargs)

hooked_transformer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:377614f5ddff670828a73b7351e6125aa6877f939184d47d27d41e9615f77b69
+size 113993767

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a460c94d9f5a51fdef98c6e0a8c84f4855b8c7d7691eed2dd16588ad2ae4e1c5
+size 100969696

modeling_interpgpt.py

@@ -0,0 +1,194 @@
+"""
+HuggingFace PreTrainedModel wrapper for InterpGPT / TaskGPT.
+
+Weights map 1:1 to the original gpt_model.TaskGPT state dict, so the same
+.pt checkpoints produced during Phase 1 load here without remapping.
+
+Usage (after upload):
+    from transformers import AutoModel, AutoTokenizer
+    model = AutoModel.from_pretrained("connaaa/interpgpt-standard-23M",
+                                      trust_remote_code=True)
+    # Or for the analysis pipeline:
+    from transformer_lens import HookedTransformer
+    hooked = HookedTransformer.from_pretrained("connaaa/interpgpt-standard-23M",
+                                               hf_model=model,
+                                               ...)
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+
+from .configuration_interpgpt import InterpGPTConfig
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d_model: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d_model))
+        self.eps = eps
+
+    def forward(self, x):
+        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x * norm * self.weight
+
+
+class RotaryPositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, max_seq_len: int = 512, base: float = 10000.0):
+        super().__init__()
+        assert d_model % 2 == 0
+        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
+        self.register_buffer("inv_freq", inv_freq)
+        t = torch.arange(max_seq_len, dtype=torch.float)
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        self.register_buffer("cos_cached", freqs.cos())
+        self.register_buffer("sin_cached", freqs.sin())
+
+    def forward(self, seq_len: int):
+        return self.cos_cached[:seq_len], self.sin_cached[:seq_len]
+
+
+def apply_rotary_emb(x, cos, sin):
+    d_half = x.shape[-1] // 2
+    x1, x2 = x[..., :d_half], x[..., d_half:]
+    cos = cos[: x.shape[2]].unsqueeze(0).unsqueeze(0)
+    sin = sin[: x.shape[2]].unsqueeze(0).unsqueeze(0)
+    out1 = x1 * cos - x2 * sin
+    out2 = x2 * cos + x1 * sin
+    return torch.cat([out1, out2], dim=-1)
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config: InterpGPTConfig):
+        super().__init__()
+        assert config.d_model % config.n_heads == 0
+        self.n_heads = config.n_heads
+        self.head_dim = config.d_model // config.n_heads
+        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=config.bias)
+        self.out_proj = nn.Linear(config.d_model, config.d_model, bias=config.bias)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.rope = RotaryPositionalEncoding(self.head_dim, config.max_seq_len)
+        mask = torch.tril(torch.ones(config.max_seq_len, config.max_seq_len))
+        self.register_buffer("causal_mask", mask.view(1, 1, config.max_seq_len, config.max_seq_len))
+
+    def forward(self, x, kv_cache=None):
+        B, T, D = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        cos, sin = self.rope(T)
+        q = apply_rotary_emb(q, cos, sin)
+        k = apply_rotary_emb(k, cos, sin)
+        if kv_cache is not None:
+            if "k" in kv_cache:
+                k = torch.cat([kv_cache["k"], k], dim=2)
+                v = torch.cat([kv_cache["v"], v], dim=2)
+            kv_cache["k"] = k
+            kv_cache["v"] = v
+        if hasattr(F, "scaled_dot_product_attention") and kv_cache is None:
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=None,
+                dropout_p=self.attn_dropout.p if self.training else 0.0,
+                is_causal=True,
+            )
+        else:
+            scale = 1.0 / math.sqrt(self.head_dim)
+            attn = torch.matmul(q, k.transpose(-2, -1)) * scale
+            T_k = k.size(2)
+            causal = self.causal_mask[:, :, T_k - T : T_k, :T_k]
+            attn = attn.masked_fill(causal == 0, float("-inf"))
+            attn = F.softmax(attn, dim=-1)
+            attn = self.attn_dropout(attn)
+            out = torch.matmul(attn, v)
+        out = out.transpose(1, 2).contiguous().view(B, T, D)
+        return self.resid_dropout(self.out_proj(out))
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: InterpGPTConfig):
+        super().__init__()
+        hidden = int(2 * config.d_ff / 3)
+        hidden = 64 * ((hidden + 63) // 64)
+        self.gate_proj = nn.Linear(config.d_model, hidden, bias=config.bias)
+        self.up_proj = nn.Linear(config.d_model, hidden, bias=config.bias)
+        self.down_proj = nn.Linear(hidden, config.d_model, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: InterpGPTConfig):
+        super().__init__()
+        self.ln1 = RMSNorm(config.d_model)
+        self.attn = CausalSelfAttention(config)
+        self.ln2 = RMSNorm(config.d_model)
+        self.ffn = FeedForward(config)
+
+    def forward(self, x, kv_cache=None):
+        x = x + self.attn(self.ln1(x), kv_cache)
+        x = x + self.ffn(self.ln2(x))
+        return x
+
+
+class InterpGPTModel(PreTrainedModel):
+    """
+    HF-wrapped InterpGPT / TaskGPT. State dict parameter names match the
+    original gpt_model.TaskGPT exactly so Phase 1 .pt checkpoints load
+    via state_dict without remapping.
+    """
+    config_class = InterpGPTConfig
+    base_model_prefix = "interpgpt"
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: InterpGPTConfig):
+        super().__init__(config)
+        self.config = config
+        self.token_embedding = nn.Embedding(config.vocab_size, config.d_model, padding_idx=config.pad_id)
+        self.drop = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
+        self.ln_final = RMSNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.token_embedding.weight
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.padding_idx is not None:
+                nn.init.zeros_(module.weight[module.padding_idx])
+
+    def forward(self, input_ids, attention_mask=None, labels=None, loss_mask=None, **kwargs):
+        B, T = input_ids.shape
+        x = self.drop(self.token_embedding(input_ids))
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_final(x)
+        logits = self.lm_head(x)
+        output = {"logits": logits}
+        if labels is not None:
+            shift_logits = logits[:, :-1].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1),
+                ignore_index=self.config.pad_id,
+                reduction="none",
+            ).view(B, T - 1)
+            if loss_mask is not None:
+                shift_mask = loss_mask[:, 1:].contiguous().float()
+                loss = (loss * shift_mask).sum() / shift_mask.sum().clamp(min=1.0)
+            else:
+                loss = loss.mean()
+            output["loss"] = loss
+        return output

pytorch_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d983e7003d6d7f5e5f28b74e06f6375a38f7cbf613616e6729eb6ab13d7358f1
+size 288797081