J space algos

gistfile1.txt

Here is the complete algorithmic pipeline extracted from the Anthropic paper, covering the training of the lens, readout, sparse J-space decomposition, and the core interventions.

---

### Algorithm 1: Train the Jacobian Lens (J_ℓ)

**Input:** Model with layers ℓ = 1 … L, corpus of prompts 𝒫, target layer = final layer (default)  
**Output:** Per-layer Jacobian matrices J_ℓ ∈ ℝ^{d_model × d_model}

```python
# h_ℓ[t] : residual stream at layer ℓ, token position t
# z[t]   : residual stream at target layer L, position t
# T      : number of token positions in the prompt

for each prompt p in corpus 𝒫:
    run forward pass; cache h_ℓ[t] for all ℓ, t

    # one backward pass per output dimension (batched in practice)
    for i in 1 … d_model:
        # inject ∂/∂z_i = 1 at EVERY target position
        grad_z = e_i ⊗ 1_T          # one-hot in dim i, broadcast over all T positions

        for each layer ℓ:
            # autodiff: gradient of summed target-layer activations w.r.t. h_ℓ
            G_ℓ = ∂( Σ_{t'=1..T} z[t'] ) / ∂h_ℓ     # shape [T, d_model]

            # average over SOURCE positions t
            J_ℓ^(p)[i, :] = mean over t of G_ℓ[t, :]

# aggregate across prompts (element-wise mean)
for each layer ℓ:
    J_ℓ = mean over prompts p of J_ℓ^(p)
```

**Mathematical definition** (§2.1):
$$J_\ell = \mathbb{E}_{t,\; t' \ge t,\; \text{prompt}} \left[ \frac{\partial h_{\text{final},t'}}{\partial h_{\ell,t}} \right]$$

---

### Algorithm 2: J-Lens Readout (Inference)

**Input:** Trained J_ℓ, residual stream activation h_ℓ, unembedding matrix W_U, normalization `norm`  
**Output:** Probability distribution over the vocabulary

```python
def lens_readout(h_ℓ, J_ℓ, W_U):
    # Project activation through the averaged Jacobian
    y = J_ℓ @ h_ℓ

    # Apply the model's pre-unembedding normalization
    y_norm = norm(y)

    # Unembed to logits and softmax
    logits = W_U @ y_norm
    probs  = softmax(logits)

    # Return ranked token list
    return argsort(probs, descending=True)
```

**Per-token probe** (§2.5): For a specific concept token *t*, read its activation as the inner product ⟨v_t, h_ℓ⟩ where v_t is the row of W_U J_ℓ corresponding to token *t*.

---

### Algorithm 3: J-Space Sparse Decomposition (Gradient Pursuit)

**Input:** Activation h_ℓ, J-lens dictionary V = {v_1 … v_{Vocab}}, sparsity budget k (typically k ≤ 25)  
**Output:** Sparse coefficients a ∈ ℝ^{Vocab}_{≥0}, J-space component h_J, non-J residual h_⊥

```python
def jspace_decompose(h, V, k):
    # Solve via gradient pursuit / matching pursuit:
    #   min || h - Σ_j a_j v_j ||_2
    #   s.t. a_j ≥ 0,  ||a||_0 ≤ k

    a = gradient_pursuit(h, dictionary=V, sparsity=k, non_negative=True)

    h_J = Σ_j a_j v_j          # J-space component
    h_⊥ = h - h_J              # non-J-space remainder

    return a, h_J, h_⊥
```

**Occupancy** (§4.2): The smallest k such that adding a (k+1)-th J-lens vector does not improve reconstruction more than adding a random direction of the same norm.

---

### Algorithm 4: Lens-Coordinate Swap (Causal Intervention)

**Input:** Source token *s*, target token *t*, activation h, swap strength α (default 1.0)  
**Output:** Patched activation h_patched

```python
def coordinate_swap(h, s, t, α=1.0):
    # Form the 2-column matrix of lens vectors
    V = [v_s  v_t]             # shape [d_model, 2]

    # Read the lens coordinates (pseudoinverse)
    c = pinv(V) @ h            # c = [c_s, c_t]

    # Swap the two entries (optionally scaled by α)
    c_swapped = [c_t, c_s] * α + [c_s, c_t] * (1 - α)
    # More precisely: σ(c) swaps the two entries, then:
    #   h_patched = h + V @ (σ(c) - c)

    delta = V @ (σ(c) - c)
    h_patched = h + delta

    return h_patched
```

**Usage:** Apply this at every token position across a band of intermediate workspace layers (e.g., L38–L92), then run the forward pass from that layer onward.

---

### Algorithm 5: J-Space Ablations & Steering

#### A. Concept-Specific Ablation
```python
def ablate_concept(h, concept_token, layers):
    v = lens_vector(concept_token)
    for h_ℓ in layers:
        # Project out the component along v
        h_ℓ = h_ℓ - (proj_v(h_ℓ))
    return h
```

#### B. Top-k J-Space Ablation (§3.5.2)
```python
def ablate_topk_jspace(h_ℓ, k=10, exclude_output_tokens=True):
    # Identify the k most strongly activated J-lens vectors
    topk = top_k_activated_lens_vectors(h_ℓ, k)

    if exclude_output_tokens:
        # Do NOT ablate tokens that appear in top-10 clean forward-pass outputs
        # (avoids corrupting imminent motor output)
        topk = filter(topk, lambda tok: tok not in clean_top10)

    # Zero out projections onto all selected vectors
    for v in topk:
        h_ℓ = h_ℓ - projection_of(h_ℓ onto v)

    return h_ℓ
```

#### C. Steering (Injection)
```python
def steer(h, concept_token, strength=α):
    v = lens_vector(concept_token)
    return h + α * v
```

---

### Algorithm 6: Broadcast-Head Ablation (§4.3.2)

**Input:** Set of J-lens "broadcast heads" ℬ (top 1% of attention heads by gain + label preservation metrics)  
**Output:** Modified forward pass with broadcast suppressed

```python
def forward_with_broadcast_ablation(x, layers, broadcast_heads ℬ):
    for ℓ in layers:
        for head h in layer ℓ:
            if h in ℬ:
                # Zero this head's output contribution
                x = x + 0 * head_output(h, x)
            else:
                x = x + head_output(h, x)
    return x
```

---

### Key Hyperparameters from the Paper

| Parameter | Typical Value | Usage |
|-----------|---------------|-------|
| Corpus size for J_ℓ | 1,000 prompts | Training the expectation |
| Sequence length | 128 tokens | Prompt sampling |
| Sparsity k | 10–25 | J-space decomposition / ablation |
| Workspace layers | L38–L92 (Sonnet 4.5) | Where interventions are applied |
| Swap strength α | 1.0 (or 2.0 for "double strength") | Coordinate swap |
| Steering strength | 0.02–0.32 | Injected-thought introspection |

---

### Quick Reference: Formal J-Space Definition (§A.8)

For a sparsity level *k* and vocabulary vectors {v_1 … v_n}, the J-space is the union of *k*-dimensional cones:

$$\mathcal{F} = \bigcup_{|S|=k} \text{span}_{\ge 0}\{v_i : i \in S\}$$

The distance from an activation *x* to the J-space is:

$$d_{\mathcal{F}}(x) = \min_{|S|=k} \|x - \Pi_S x\|$$

where Π_S is orthogonal projection onto the cone spanned by the chosen *k* vectors. The minimizing projection is the **J-space component** used in all interventions above.