wassname/AntiPaSTO
1
0
Fork 0
mirror of https://github.com/wassname/AntiPaSTO.git synced 2026-09-10 11:40:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

simplify

Browse Source
This commit is contained in:
wassname
2026-01-13 08:43:04 +08:00
parent af60db82f9
commit da4dda8ab7
2 changed files with 36 additions and 972 deletions
Download Patch File Download Diff File Expand all files Collapse all files
antipasto/peft_utils/layer_selection.py
+23 -299
View File
@@ -1,28 +1,22 @@
#!/usr/bin/env python3
"""Centralized layer selection logic for AntiPaSTO training.
FLOW OVERVIEW (v2.2+)
=====================
1. compute_simple_layer_selection() [RECOMMENDED]:
- Computes SVD for all linear layers (needed for adapter init anyway)
- Selects layers uniformly across valid depth range
- Uses top-r singular values for dimension selection (no gradient)
- Computes weight-only subspaces (write, write_x_notlogits)
- No backward pass = no OOM on large models (12B+)
SIMPLIFIED FOR PUBLICATION (2026-01-13):
Only 3 loss_subspace types are supported: taskdiff_x_suppressed_x_write (default), write, taskdiff.
Deprecated gradient-based layer selection has been removed.
2. compute_gradient_layer_selection() [DEPRECATED]:
- Expensive backward pass for gradient-based ranking
- Ablations show it doesn't improve over simple selection
- Kept for research/debugging purposes
FLOW:
1. compute_simple_layer_selection():
- Computes SVD for all linear layers (needed for adapter init)
- Selects layers uniformly across valid depth range
- Computes subspaces: write, taskdiff, taskdiff_x_suppressed_x_write
Key functions:
- compute_simple_layer_selection(): Uniform layer selection, top-S dims, weight-only subspaces
- compute_gradient_layer_selection(): Gradient-based selection (deprecated, OOMs on 12B+)
- compute_simple_layer_selection(): Main entry point for layer/subspace selection
- find_linear_layers(): Discover all linear modules in model
- resolve_target_modules(): Expand "residual-writers" etc. to concrete module lists
Subspace operations (compute_write_subspace, compute_write_x_notlogits, etc.)
are in antipasto/peft_utils/subspaces.py
Subspace operations are in antipasto/peft_utils/subspaces.py
"""
import re
import pandas as pd
@@ -41,23 +35,11 @@ from torch.utils.data import DataLoader, Subset
from transformers import DataCollatorWithPadding
import gc
from antipasto.peft_utils.subspaces import (
compute_lm_head_subspace,
compute_lm_head_svd,
compute_suppressed_from_hidden_states,
compute_churn_from_hidden_states,
compute_churn_constructive_from_hidden_states,
compute_churn_suppressive_from_hidden_states,
compute_task_diff_from_hidden_states,
compute_task_diff_constructive_from_hidden_states,
compute_task_read_subspace,
compute_task_lm_head_subspace,
compute_task_wnr_subspace,
compute_module_subspace_from_svds,
compute_write_not_read_subspace,
compute_stenographic_subspace,
compute_write_x_notlogits_subspace,
compute_logits_tail_subspace,
compute_taskdiff_x_write_x_notlogits_subspace,
find_write_modules,
find_read_modules,
approx_intersection,
@@ -944,304 +926,46 @@ def compute_simple_layer_selection(
S=lm_head_S_full[:INTERMEDIATE_SUBSPACE_RANK].to(device=device, dtype=dtype).detach(),
)
if lm_head_sub is not None:
subspaces.set('logits_read', lm_head_sub) # Full Subspace with S
# write_x_notlogits = write projected into notlogits (compound: needs full geometry)
if write_subspace is not None and lm_head_sub is not None:
hfl_sub = compute_write_x_notlogits_subspace(
write_subspace=write_subspace,
lm_head_subspace=lm_head_sub,
top_k=INTERMEDIATE_SUBSPACE_RANK, # Full geometry for subtraction
)
subspaces.set('write_x_notlogits', hfl_sub) # Full Subspace
logger.info(f"write_x_notlogits subspace: rank={hfl_sub.V.shape[1]}")
# Read subspace from read modules (q_proj, k_proj, etc.) row spaces
read_modules = find_read_modules(model)
read_subspace = compute_module_subspace_from_svds(
layer_svds=layer_svd_cpu,
layer_info=layer_info_full,
module_filter=read_modules,
use_column_space=False, # Row space = read directions
top_k=INTERMEDIATE_SUBSPACE_RANK,
device=device,
dtype=dtype,
name="read",
)
if read_subspace is not None:
subspaces.set('read', read_subspace) # Full Subspace
logger.info(f"read subspace: rank={read_subspace.V.shape[1]}")
# NEW: Specific read subspaces (Query, Key, Value)
for name, suffix in [('query_read', 'q_proj'), ('key_read', 'k_proj'), ('value_read', 'v_proj')]:
modules = [m for m in read_modules if suffix in m]
if modules:
sub = compute_module_subspace_from_svds(
layer_svds=layer_svd_cpu,
layer_info=layer_info_full,
module_filter=modules,
use_column_space=False, # Row space
top_k=INTERMEDIATE_SUBSPACE_RANK,
device=device,
dtype=dtype,
name=name,
)
if sub is not None:
subspaces.set(name, sub)
# Attention Read (Union of Q, K, V)
attn_read_modules = [m for m in read_modules if any(s in m for s in ['q_proj', 'k_proj', 'v_proj'])]
if attn_read_modules:
attn_read_subspace = compute_module_subspace_from_svds(
layer_svds=layer_svd_cpu,
layer_info=layer_info_full,
module_filter=attn_read_modules,
use_column_space=False, # Row space = read directions
top_k=INTERMEDIATE_SUBSPACE_RANK,
device=device,
dtype=dtype,
name="attn_read",
)
# NEW: Attention Sink = Write - Attention_Read
if write_subspace is not None and attn_read_subspace is not None:
attn_sink = compute_write_not_read_subspace(
write_subspace=write_subspace,
read_subspace=attn_read_subspace,
lm_head_subspace=None, # Don't subtract lm_head for pure attention sink
top_k=INTERMEDIATE_SUBSPACE_RANK,
)
subspaces.set('attention_sink', attn_sink)
logger.info(f"attention_sink subspace: intermediate={attn_sink.V.shape[1]}, stored={top_k}")
# NEW: Communication Channel = Write & Read
if write_subspace is not None and read_subspace is not None:
comm_channel = approx_intersection(write_subspace, read_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('communication_channel', comm_channel)
logger.info(f"communication_channel subspace: intermediate={comm_channel.V.shape[1]}, stored={top_k}")
# write_not_read = write - read - lm_head (what's written but ignored)
if write_subspace is not None and read_subspace is not None:
wnr_sub = compute_write_not_read_subspace(
write_subspace=write_subspace,
read_subspace=read_subspace,
lm_head_subspace=lm_head_sub, # Also subtract lm_head if available
top_k=INTERMEDIATE_SUBSPACE_RANK,
)
subspaces.set('write_not_read', wnr_sub)
logger.info(f"write_not_read subspace: intermediate={wnr_sub.V.shape[1]}, stored={top_k}")
subspaces.set('logits_read', lm_head_sub)
# =========================================================================
# WANDA_X_LOGITS_NULL SUBSPACE
# =========================================================================
if loss_subspace == 'wanda_x_notlogits':
if hs_stacked is None:
raise ValueError("wanda_x_notlogits requires hidden states")
logger.info("Computing wanda_x_notlogits subspace...")
# Compute lm_head SVD to get S and Vh
W = model.lm_head.weight.data.float().cpu()
_, S_lm_head, Vh_lm_head = torch.linalg.svd(W, full_matrices=False)
# logits_tail uses its own internal PCA - pass intermediate rank
# for full geometry, then truncate at storage
active_null_sub = compute_logits_tail_subspace(
hidden_states=hs_stacked,
lm_head_S=S_lm_head,
lm_head_Vh=Vh_lm_head,
top_k=INTERMEDIATE_SUBSPACE_RANK
)
subspaces.set('wanda_x_notlogits', active_null_sub)
logger.info(f"wanda_x_notlogits subspace: intermediate={active_null_sub.V.shape[1]}, stored={top_k}")
# Random subspace (sanity baseline)
if loss_subspace == 'random':
if write_subspace is not None:
d_model = write_subspace.V.shape[0]
elif lm_head_sub is not None:
d_model = lm_head_sub.V.shape[0]
else:
# Fall back to any SVD to infer d_model (input dim for residual-connected linears)
any_path = next(iter(layer_svd_cpu.keys()))
_, _, any_Vh = layer_svd_cpu[any_path]
d_model = any_Vh.shape[1]
gen_random = torch.Generator(device=device)
gen_random.manual_seed(_stable_u32(f"{seed}:loss_subspace:random:{d_model}:{top_k}"))
random_basis = torch.randn(d_model, top_k, device=device, dtype=dtype, generator=gen_random)
random_basis = torch.linalg.qr(random_basis.float())[0].to(dtype) # Orthonormalize
subspaces.set('random', random_basis)
logger.info(f"random subspace: shape={random_basis.shape}")
# =========================================================================
# ACTIVATION-BASED SUBSPACES (uses hidden states collected earlier)
# Always computed when hidden states are available (dataset_pt was provided)
# ACTIVATION-BASED SUBSPACES (simplified for publication)
# Only computes subspaces needed for 3 supported loss_subspace types:
# - taskdiff, taskdiff_x_suppressed_x_write (default)
# =========================================================================
if hs_stacked is not None:
logger.info(f"Computing activation-based subspaces (taskdiff, churn, suppressed, etc.)...")
logger.info(f"Computing activation-based subspaces (taskdiff, suppressed)...")
# Compute taskdiff subspace - use INTERMEDIATE rank for full geometry
# Compute taskdiff subspace - PCA on cho-rej difference
task_diff_subspace = compute_task_diff_from_hidden_states(
hidden_states=hs_stacked,
top_k=INTERMEDIATE_SUBSPACE_RANK,
layer_frac=loss_hs_frac_for_task,
)
subspaces.set('taskdiff', task_diff_subspace) # Full Subspace with S for energy thresholding
# taskdiff_write: per-layer contributions that differ between cho/rej
taskdiff_write_subspace = compute_task_diff_from_hidden_states(
hidden_states=hs_stacked,
top_k=INTERMEDIATE_SUBSPACE_RANK,
layer_frac=loss_hs_frac_for_task,
use_layer_diffs=True,
)
subspaces.set('taskdiff_write', taskdiff_write_subspace)
subspaces.set('taskdiff', task_diff_subspace)
# Compute suppressed subspace (from layer diffs) - use INTERMEDIATE rank
# Compute suppressed subspace (written but erased by later layers)
suppressed_subspace = compute_suppressed_from_hidden_states(
hidden_states=hs_stacked,
lm_head_subspace=lm_head_sub,
top_k=INTERMEDIATE_SUBSPACE_RANK,
)
subspaces.set('suppressed', suppressed_subspace) # Full Subspace with S for energy thresholding
subspaces.set('suppressed', suppressed_subspace)
# Compute churn subspace - use INTERMEDIATE rank
churn_subspace = compute_churn_from_hidden_states(
hidden_states=hs_stacked,
top_k=INTERMEDIATE_SUBSPACE_RANK,
)
subspaces.set('churn', churn_subspace) # Full Subspace with S for energy thresholding
# taskdiff_x_suppressed = taskdiff ∩ suppressed (compound: needs full geometry)
# taskdiff_x_suppressed = taskdiff ∩ suppressed (stenographic signal)
steno_subspace = compute_stenographic_subspace(
task_diff_subspace=task_diff_subspace,
suppressed_subspace=suppressed_subspace,
top_k=INTERMEDIATE_SUBSPACE_RANK, # Full geometry for intersection
)
subspaces.set('taskdiff_x_suppressed', steno_subspace) # Store full Subspace for energy thresholding
# NEW: Prediction Suppression = Suppressed & Read
if suppressed_subspace is not None and read_subspace is not None:
pred_supp = approx_intersection(suppressed_subspace, read_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('prediction_suppression', pred_supp)
logger.info(f"prediction_suppression subspace: intermediate={pred_supp.V.shape[1]}, stored={top_k}")
# Compound subspaces: taskdiff ∩ X (all need full geometry for intersection)
if hfl_sub is not None:
task_intersect_hfl = approx_intersection(task_diff_subspace, hfl_sub, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_write_x_notlogits', task_intersect_hfl)
if write_subspace is not None:
task_intersect_write = approx_intersection(task_diff_subspace, write_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_write', task_intersect_write)
task_intersect_churn = approx_intersection(task_diff_subspace, churn_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_churn', task_intersect_churn)
task_intersect_steno = approx_intersection(task_diff_subspace, steno_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_taskdiff_x_suppressed', task_intersect_steno)
# Churn variants (constructive = magnitude increase, suppressive = magnitude decrease)
churn_constructive = compute_churn_constructive_from_hidden_states(
hidden_states=hs_stacked,
top_k=INTERMEDIATE_SUBSPACE_RANK,
)
subspaces.set('churn_constructive', churn_constructive)
churn_suppressive = compute_churn_suppressive_from_hidden_states(
hidden_states=hs_stacked,
top_k=INTERMEDIATE_SUBSPACE_RANK,
)
subspaces.set('churn_suppressive', churn_suppressive)
# task_diff_constructive = directions where task magnitude INCREASES
task_diff_constructive = compute_task_diff_constructive_from_hidden_states(
hidden_states=hs_stacked,
top_k=INTERMEDIATE_SUBSPACE_RANK,
layer_range=(min_adapter_layer_frac, loss_layer_frac),
)
subspaces.set('taskdiff_constructive', task_diff_constructive)
# Task compound subspaces with specific weight subspaces
if write_subspace is not None:
task_write = approx_intersection(task_diff_subspace, write_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_write', task_write)
if read_subspace is not None:
# task ∩ read: task-discriminative directions that are read by attention/MLP inputs
task_read = approx_intersection(task_diff_subspace, read_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_read', task_read)
# task ∩ lm_head: task directions that affect output logits
if lm_head_sub is not None:
task_lm_head = compute_task_lm_head_subspace(
task_diff_subspace=task_diff_subspace,
lm_head_subspace=lm_head_sub,
top_k=INTERMEDIATE_SUBSPACE_RANK,
)
subspaces.set('taskdiff_x_logits_read', task_lm_head)
if wnr_sub is not None:
task_wnr = compute_task_wnr_subspace(
task_diff_subspace=task_diff_subspace,
write_not_read_subspace=wnr_sub,
top_k=INTERMEDIATE_SUBSPACE_RANK,
)
subspaces.set('taskdiff_x_write_not_read', task_wnr)
# Additional compound intersections for sweeps
task_intersect_churn_constructive = approx_intersection(task_diff_subspace, churn_constructive, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_churn_constructive', task_intersect_churn_constructive)
if wnr_sub is not None and write_subspace is not None:
taskdiff_write_intersect_wnr = approx_intersection(taskdiff_write_subspace, wnr_sub, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_write_x_write_not_read', taskdiff_write_intersect_wnr)
taskdiff_write_intersect_suppressed = approx_intersection(taskdiff_write_subspace, suppressed_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_write_x_suppressed', taskdiff_write_intersect_suppressed)
taskdiff_write_intersect_churn = approx_intersection(taskdiff_write_subspace, churn_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_write_x_churn', taskdiff_write_intersect_churn)
# task_constructive_intersect_* (task_diff_constructive ∩ X)
if hfl_sub is not None:
task_constructive_intersect_hfl = approx_intersection(task_diff_constructive, hfl_sub, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_constructive_x_write_x_notlogits', task_constructive_intersect_hfl)
task_constructive_intersect_suppressed = approx_intersection(task_diff_constructive, suppressed_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_constructive_x_suppressed', task_constructive_intersect_suppressed)
task_constructive_intersect_churn = approx_intersection(task_diff_constructive, churn_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_constructive_x_churn', task_constructive_intersect_churn)
task_constructive_intersect_churn_constructive = approx_intersection(task_diff_constructive, churn_constructive, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_constructive_x_churn_constructive', task_constructive_intersect_churn_constructive)
# taskdiff_x_suppressed_x_* (taskdiff_x_suppressed ∩ X)
steno_intersect_churn = approx_intersection(steno_subspace, churn_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_suppressed_x_churn', steno_intersect_churn)
subspaces.set('taskdiff_x_suppressed', steno_subspace)
# taskdiff_x_suppressed_x_write = steno ∩ write (default loss subspace)
if write_subspace is not None:
steno_intersect_write = approx_intersection(steno_subspace, write_subspace, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_suppressed_x_write', steno_intersect_write)
if hfl_sub is not None:
steno_intersect_hfl = approx_intersection(steno_subspace, hfl_sub, top_k=INTERMEDIATE_SUBSPACE_RANK)
subspaces.set('taskdiff_x_suppressed_x_write_x_notlogits', steno_intersect_hfl)
# taskdiff_write_x_notlogits: task-discriminative directions in write ∩ notlogits
# (like write_x_notlogits but weighted by cho-rej difference)
if write_subspace is not None and lm_head_sub is not None:
taskdiff_x_write_x_notlogits = compute_taskdiff_x_write_x_notlogits_subspace(
hidden_states=hs_stacked,
write_subspace=write_subspace,
lm_head_S=lm_head_S_full,
lm_head_Vh=lm_head_Vh_full,
top_k=INTERMEDIATE_SUBSPACE_RANK,
layer_frac=loss_layer_frac,
)
subspaces.set('taskdiff_write_x_notlogits', taskdiff_x_write_x_notlogits)
logger.info(f"Activation-based subspaces computed (intermediate={INTERMEDIATE_SUBSPACE_RANK}, stored={top_k}): {[k for k in subspaces.keys() if 'task' in k or 'steno' in k or 'churn' in k]}")
logger.info(f"Activation-based subspaces computed: taskdiff, suppressed, taskdiff_x_suppressed, taskdiff_x_suppressed_x_write")
# Cleanup hidden states if collected
if hs_stacked is not None:
antipasto/peft_utils/subspaces.py
+13 -673
View File
@@ -18,6 +18,19 @@ Naming conventions for subspace functions:
For geometric intuition and taxonomy of named subspaces, see docs/steering_methods.qmd.
All bases are detached (frozen) to prevent gradient hacking.
SIMPLIFICATION NOTE (2026-01-13):
This file was simplified for publication. Many experimental subspace types were removed.
Only the following are supported:
- taskdiff_x_suppressed_x_write (default)
- write
- taskdiff
See git history (pre-2026-01-13) for removed experimental types:
- write_not_read, stenographic, write_x_notlogits, logits_tail
- taskdiff_x_write_x_notlogits, task_read, task_lm_head, task_wnr
- churn variants (constructive, suppressive), taskdiff_constructive
- lm_head, embed subspaces
"""
from __future__ import annotations
@@ -647,37 +660,6 @@ def compute_suppressed_from_hidden_states(
return Subspace(V_supp, name="suppressed", S=S_supp)
# ============================================================================
# Legacy subspace computation (kept for compatibility)
# ============================================================================
def compute_lm_head_subspace(model: nn.Module, top_k: int = 256) -> Subspace:
"""Compute subspace read by lm_head (directions that affect output logits).
Uses right singular vectors (V) of lm_head.weight since it reads from residual.
lm_head computes logits = h @ W.T, so it reads directions in row-space of W.
Args:
model: Model with lm_head
top_k: Number of components
Returns:
Subspace of directions lm_head reads
"""
device = next(model.parameters()).device
dtype = next(model.parameters()).dtype
# SVD: W = U @ S @ Vh, row-space = span of Vh rows = right singular vectors
S, Vh = compute_lm_head_svd(model)
V_read: Float[Tensor, "d_model top_k"] = Vh[:top_k, :].T # transpose: [d_model, top_k]
S_read = S[:top_k].to(dtype).to(device).detach()
V_read = V_read.to(dtype).to(device).detach()
logger.debug(f"logits_read subspace: rank={V_read.shape[1]}")
return Subspace(V_read, name="logits_read", S=S_read)
def compute_lm_head_svd(model: nn.Module) -> tuple[Tensor, Tensor]:
"""Return (S, Vh) for lm_head.weight SVD.
@@ -694,68 +676,6 @@ def compute_lm_head_svd(model: nn.Module) -> tuple[Tensor, Tensor]:
return S, Vh
def compute_embed_subspace(model: nn.Module, top_k: int = 256) -> Subspace:
"""Compute subspace written by embedding layer.
Uses left singular vectors (U) of embed_tokens.weight since it writes to residual.
Args:
model: Model with embed_tokens
top_k: Number of components
Returns:
Subspace of directions embedding writes
"""
device = next(model.parameters()).device
dtype = next(model.parameters()).dtype
# embed_tokens.weight: [vocab_size, d_model], output is row-indexed
# Column space = write directions
W = model.model.embed_tokens.weight.data # [vocab, d_model]
# Column space via transpose
U, S, _ = torch.linalg.svd(W.T.float().cpu(), full_matrices=False)
V_write = U[:, :top_k] # [d_model, top_k]
V_write = V_write.to(dtype).to(device).detach()
logger.info(f"Embed write subspace: rank={V_write.shape[1]}")
return Subspace(V_write, name="embed_write")
def compute_write_not_read_subspace(
write_subspace: Subspace,
read_subspace: Subspace,
lm_head_subspace: Optional[Subspace] = None,
top_k: int =256,
) -> Subspace:
"""Compute Write-Not-Read subspace: directions written but not read.
Notation: WnR = Write_perp_Read = Π_{Read^⊥}(Write).
If `lm_head_subspace` is provided, also subtract directions readable by
the lm_head (since those are "read" at the output interface).
Args:
write_subspace: Subspace of write directions
read_subspace: Subspace of read directions
top_k: Number of components
Returns:
Subspace of directions written but ignored by reading layers
"""
wnr = project_subspace_into_perp(write_subspace, read_subspace)
if lm_head_subspace is not None:
wnr = project_subspace_into_perp(wnr, lm_head_subspace)
if wnr.rank > top_k:
wnr = Subspace(wnr.V[:, :top_k], name="write_not_read")
else:
wnr.name = "write_not_read"
return wnr
def compute_stenographic_subspace(
task_diff_subspace: Subspace,
suppressed_subspace: Subspace,
@@ -778,504 +698,6 @@ def compute_stenographic_subspace(
return steno
def compute_write_x_notlogits_subspace(
write_subspace: Subspace,
lm_head_subspace: Subspace,
top_k: int =256,
) -> Subspace:
"""Compute write_x_notlogits: write projected into (logits_read)^perp.
Notation: write_x_notlogits = Write_perp_logits_read = Π_{(logits_read)^⊥}(Write).
In code this uses project_subspace_into_perp(write, logits), which performs an
orthogonal-complement projection (see project_bases_into_perp docstring), not a set
difference.
These directions are written to residual by model layers but don't affect
output logits (lm_head can't read them). Simpler than write_not_read since
it ignores layer-to-layer reads.
Note it includes write to avoid token embeddings that prepopulate the residual stream
Args:
write_subspace: Subspace of write directions
lm_head_subspace: Subspace readable by lm_head
top_k: Number of components
Returns:
Subspace of directions hidden from final output
"""
hfl = project_subspace_into_perp(write_subspace, lm_head_subspace)
if hfl.rank > top_k:
hfl = Subspace(hfl.V[:, :top_k], name="write_x_notlogits")
else:
hfl.name = "write_x_notlogits"
return hfl
def compute_logits_tail_subspace(
hidden_states: Float[Tensor, "batch n_layers_plus1 d_model"],
lm_head_S: Float[Tensor, "rank"],
lm_head_Vh: Float[Tensor, "rank d_model"],
top_k: int = 64,
layer_range: Optional[tuple] = None,
null_frac: float = 0.5,
) -> Subspace:
"""Compute wanda_x_notlogits subspace: tail lm_head singular dirs weighted by activation.
Like write_x_notlogits but empirical: uses actual activations to weight directions.
Method (WANDA-inspired):
1. Take bottom `null_frac` of lm_head singular directions (low S = low output gain)
2. Project hidden states into this tail subspace
3. Weight each direction by activation magnitude (WANDA: ||X||_2 per direction)
4. PCA on weighted projections to find most-used directions within tail space
This differs from write_x_notlogits (static weight subtraction) by incorporating
which directions are actually used, not just which could theoretically be hidden.
Args:
hidden_states: [batch, n_layers+1, d_model] from model output
lm_head_S: [rank] singular values of lm_head (descending order from SVD)
lm_head_Vh: [rank, d_model] right singular vectors (rows are basis vectors)
top_k: Number of components to return
layer_range: (start_frac, end_frac) for which layers to use (default 0.3-0.8)
null_frac: Fraction of bottom singular directions to use (default 0.5)
Returns:
Subspace of actively-used low-gain directions
"""
device = hidden_states.device
dtype = hidden_states.dtype
d_model = hidden_states.shape[-1]
n_layers_plus1 = hidden_states.shape[1]
n_layers = n_layers_plus1 - 1
if layer_range is None:
layer_range = (0.3, 0.8)
start_idx = max(1, int(layer_range[0] * n_layers))
end_idx = min(n_layers, int(layer_range[1] * n_layers))
# Get relevant hidden states [batch, selected_layers, d_model]
hs_selected = hidden_states[:, start_idx:end_idx, :]
# Take bottom null_frac of singular directions (low S = lm_head ignores)
rank = lm_head_Vh.shape[0]
null_start = int((1 - null_frac) * rank)
null_rank = rank - null_start
if null_rank < top_k:
logger.warning(f"null_frac={null_frac} gives {null_rank} dims < top_k={top_k}. Expanding.")
null_start = max(0, rank - top_k * 2)
null_rank = rank - null_start
# V_tail: [d_model, null_rank] - bottom singular vectors
V_tail = lm_head_Vh[null_start:, :].T.to(device).to(dtype) # [d_model, null_rank]
S_tail = lm_head_S[null_start:].to(device).float() # [null_rank]
# Project hidden states into tail subspace
hs_flat = hs_selected.reshape(-1, d_model).float()
z = hs_flat @ V_tail.float() # [n, null_rank]
# WANDA-style: weight by activation magnitude (L2 norm per direction)
# ||X_j||_2 = sqrt(sum_i x_ij^2), captures total energy in each direction
activation_norm = z.norm(dim=0) # [null_rank] - L2 norm across samples
# Weight projections by activation norm
z_weighted = z * activation_norm # [n, null_rank]
# PCA on weighted projections
z_centered = z_weighted - z_weighted.mean(dim=0, keepdim=True)
_, S_pca, Vh_pca = torch.linalg.svd(z_centered, full_matrices=False)
# Top-k directions in tail basis
k = min(top_k, Vh_pca.shape[0])
U_top = Vh_pca[:k, :] # [k, null_rank]
# Map back to residual basis: [k, null_rank] @ [null_rank, d_model] -> [k, d_model]
V_result = (U_top @ V_tail.T.float()).T # [d_model, k]
# Orthonormalize
V_result, _ = torch.linalg.qr(V_result)
V_result = V_result[:, :k].to(dtype).to(device).detach()
explained_var = (S_pca[:k] ** 2).sum() / ((S_pca ** 2).sum() + 1e-8)
act_range = f"{activation_norm.min():.2f}-{activation_norm.max():.2f}"
s_range = f"{S_tail.min():.2e}-{S_tail.max():.2e}"
logger.info(f"wanda_x_notlogits subspace: rank={V_result.shape[1]}, null_dims={null_rank}, "
f"explained_var={explained_var:.1%}, activation_range={act_range}, S_range={s_range}")
log_topk_explained_variance(S_pca ** 2, "wanda_x_notlogits") # squared because we used variance formula
return Subspace(V_result, name="wanda_x_notlogits")
def compute_taskdiff_x_write_x_notlogits_subspace(
hidden_states: Float[Tensor, "batch n_layers_plus1 d_model"],
write_subspace: "Subspace",
lm_head_S: Float[Tensor, "rank"],
lm_head_Vh: Float[Tensor, "rank d_model"],
top_k: int = 64,
layer_frac: float = 0.7,
null_frac: float = 0.5,
) -> "Subspace":
"""Task-discriminative directions in write ∩ lm_head_null.
Finds directions that are:
1. Writable (in column space of o_proj/down_proj)
2. Hidden from lm_head (in bottom singular vectors of lm_head)
3. Task-discriminative (high cho-rej difference magnitude)
Unlike `write_x_notlogits` (weight-only), this uses cho-rej activations
to find WHICH hidden directions carry task-relevant signal.
Unlike `logits_tail` (sample-specific), this weights by cho-rej
DIFFERENCE, not total activation magnitude.
Args:
hidden_states: [batch, n_layers+1, d_model] from contrastive pairs
Assumes batch dimension alternates cho/rej: [cho_0, rej_0, cho_1, rej_1, ...]
write_subspace: Subspace of write directions (from compute_write_subspace)
lm_head_S: [rank] singular values of lm_head (descending)
lm_head_Vh: [rank, d_model] right singular vectors
top_k: Number of components to return
layer_frac: Which layer to use (fraction of total layers)
null_frac: Fraction of bottom singular vectors to use as "null" (default 0.5)
Returns:
Subspace of task-discriminative write-lm_null directions
"""
device = hidden_states.device
dtype = hidden_states.dtype
batch, n_layers_plus1, d_model = hidden_states.shape
n_layers = n_layers_plus1 - 1
# Get layer hidden states
layer_idx = int(layer_frac * n_layers)
hs = hidden_states[:, layer_idx, :].float() # [batch, d]
# Split cho/rej (assumes alternating)
hs_cho = hs[0::2] # [n_pairs, d]
hs_rej = hs[1::2] # [n_pairs, d]
diff = hs_cho - hs_rej # [n_pairs, d]
# Step 1: Get lm_head null space (bottom singular vectors = low output gain)
rank = lm_head_Vh.shape[0]
null_start = int((1 - null_frac) * rank)
V_lm_null = lm_head_Vh[null_start:, :].T.to(device).float() # [d, null_rank]
# Step 2: Intersect with write space
V_write = write_subspace.V.to(device).float() # [d, write_rank]
V_write_lmnull, _ = approx_intersection_bases(V_write, V_lm_null, top_k=256) # [d, intersect_rank]
if V_write_lmnull.shape[1] < 2:
logger.warning(f"taskdiff_x_write_x_notlogits: write ∩ lm_null intersection too small ({V_write_lmnull.shape[1]}), using write only")
V_write_lmnull = V_write
# Step 3: Project differences into write ∩ lm_null
z_diff = diff @ V_write_lmnull # [n_pairs, intersect_rank]
# Step 4: Weight by task-discriminative magnitude (mean absolute difference)
task_weight = z_diff.abs().mean(dim=0) # [intersect_rank]
# Step 5: PCA on weighted projections to find most task-discriminative directions
z_weighted = z_diff * task_weight
z_centered = z_weighted - z_weighted.mean(dim=0, keepdim=True)
_, S_pca, Vh_pca = torch.linalg.svd(z_centered, full_matrices=False)
# Top-k directions in intersection basis
k = min(top_k, Vh_pca.shape[0])
U_top = Vh_pca[:k, :] # [k, intersect_rank]
# Map back to residual basis: [k, intersect_rank] @ [intersect_rank, d] -> [k, d]
V_result = (U_top @ V_write_lmnull.T).T # [d, k]
# Orthonormalize
V_result = orthonormalize(V_result).to(dtype).to(device).detach()
explained_var = (S_pca[:k] ** 2).sum() / ((S_pca ** 2).sum() + 1e-8)
intersect_rank = V_write_lmnull.shape[1]
weight_range = f"{task_weight.min():.2f}-{task_weight.max():.2f}"
logger.info(f"taskdiff_write_x_notlogits subspace: rank={V_result.shape[1]}, intersect_rank={intersect_rank}, "
f"explained_var={explained_var:.1%}, task_weight_range={weight_range}")
log_topk_explained_variance(S_pca ** 2, "taskdiff_write_x_notlogits")
return Subspace(V_result, name="taskdiff_write_x_notlogits")
# ============================================================================
# Subspace computation from precomputed SVDs (used by layer_selection.py)
# ============================================================================
def compute_churn_from_hidden_states(
hidden_states: Float[Tensor, "batch n_layers_plus1 d_model"],
top_k: int =256,
) -> Subspace:
"""Compute churn subspace: PCA of layer-to-layer changes.
Churn captures "active computation lanes" - directions where layers
add and remove energy during processing.
Args:
hidden_states: [batch, n_layers+1, d_model] - all layer outputs
top_k: Number of components to keep
Returns:
Subspace of high-churn directions
"""
device = hidden_states.device
dtype = hidden_states.dtype
d_model = hidden_states.shape[-1]
# Layer diffs: [batch, n_layers, d_model]
layer_diffs: Float[Tensor, "batch n_layers d"] = hidden_states[:, 1:, :] - hidden_states[:, :-1, :]
layer_diffs_flat: Float[Tensor, "n d"] = layer_diffs.reshape(-1, d_model).float()
# PCA of layer diffs. normalize_samples=False: layer diffs are already comparable
# (same scale within a model), and we want magnitude-weighted to capture where
# most computation happens.
sub = pca_subspace(
layer_diffs_flat,
top_k=top_k,
normalize_samples=False,
name="churn",
device=device,
dtype=dtype,
)
return sub
def compute_churn_constructive_from_hidden_states(
hidden_states: Float[Tensor, "batch n_layers_plus1 d_model"],
top_k: int =256,
layer_range: Optional[tuple] = None,
) -> Subspace:
"""Compute constructive churn: directions where magnitude INCREASES across layers.
Standard churn is unsigned (PCA of layer diffs). This variant filters to directions
where the residual stream is actively BUILDING signal (amplifying), not erasing it.
Method: For each churn PC, compute whether ||h @ v||^2 increases from early to late layers.
Keep only PCs where slope > 0 (magnitude growing).
Args:
hidden_states: [batch, n_layers+1, d_model] - all layer outputs
top_k: Number of components to keep
layer_range: Optional (start_frac, end_frac) for slope computation (default 0.2-0.8)
Returns:
Subspace of constructive (amplifying) churn directions
"""
device = hidden_states.device
dtype = hidden_states.dtype
d_model = hidden_states.shape[-1]
n_layers_plus1 = hidden_states.shape[1]
n_layers = n_layers_plus1 - 1
if layer_range is None:
layer_range = (0.2, 0.8)
start_idx = max(1, int(layer_range[0] * n_layers))
end_idx = min(n_layers, int(layer_range[1] * n_layers))
# First compute regular churn PCs
layer_diffs: Float[Tensor, "batch n_layers d"] = hidden_states[:, 1:, :] - hidden_states[:, :-1, :]
layer_diffs_flat: Float[Tensor, "n d"] = layer_diffs.reshape(-1, d_model).float()
layer_diffs_centered = layer_diffs_flat - layer_diffs_flat.mean(dim=0, keepdim=True)
_, S, Vh = torch.linalg.svd(layer_diffs_centered, full_matrices=False)
# Get more PCs than we need to filter
n_candidates = min(top_k * 3, Vh.shape[0])
V_candidates: Float[Tensor, "d k"] = Vh[:n_candidates, :].T # [d_model, n_candidates]
# For each PC, compute magnitude trend across layers
# Project hidden states onto each PC: [batch, n_layers+1, n_candidates]
proj_mag_sq = (hidden_states.float() @ V_candidates) ** 2 # [batch, n_layers+1, n_candidates]
# Compute slope via early vs late layer magnitude.
# We average over a 3-layer window at each endpoint for noise reduction.
# The "constructive" signal is (late_mag - early_mag) > 0, meaning
# magnitude in this PC direction is INCREASING through the network.
# Window size 3 is a tradeoff: smaller = more sensitive but noisier.
early_mag = proj_mag_sq[:, start_idx:start_idx+3, :].mean(dim=(0, 1)) # [n_candidates]
late_mag = proj_mag_sq[:, end_idx-3:end_idx, :].mean(dim=(0, 1)) # [n_candidates]
# Constructive = late > early (magnitude increasing)
mag_slope = late_mag - early_mag # positive = constructive
# Select top-k by constructiveness (positive slope), sorted by magnitude
constructive_mask = mag_slope > 0
if constructive_mask.sum() < top_k:
# Fallback: take all with positive slope, fill with least negative
logger.warning(f"Only {constructive_mask.sum()} constructive PCs found, taking {top_k} least suppressive")
sorted_indices = torch.argsort(mag_slope, descending=True)[:top_k]
else:
# Among constructive, sort by explained variance (S) and take top-k
constructive_indices = torch.where(constructive_mask)[0]
# Weight by both constructiveness and variance explained
scores = mag_slope[constructive_indices] * S[constructive_indices]
sorted_by_score = torch.argsort(scores, descending=True)[:top_k]
sorted_indices = constructive_indices[sorted_by_score]
V_constructive: Float[Tensor, "d k"] = V_candidates[:, sorted_indices].to(dtype).to(device).detach()
n_positive = (mag_slope[sorted_indices] > 0).sum().item()
logger.info(f"Churn_constructive subspace: rank={V_constructive.shape[1]}, {n_positive}/{top_k} strictly constructive")
return Subspace(V_constructive, name="churn_constructive")
def compute_churn_suppressive_from_hidden_states(
hidden_states: Float[Tensor, "batch n_layers_plus1 d_model"],
top_k: int =256,
layer_range: Optional[tuple] = None,
) -> Subspace:
"""Compute suppressive churn: directions where magnitude DECREASES across layers.
Complement to constructive churn. These are directions the model is actively
ERASING or damping during processing. Steering these could fight the model's flow.
Args:
hidden_states: [batch, n_layers+1, d_model] - all layer outputs
top_k: Number of components to keep
layer_range: Optional (start_frac, end_frac) for slope computation (default 0.2-0.8)
Returns:
Subspace of suppressive (erasing) churn directions
"""
device = hidden_states.device
dtype = hidden_states.dtype
d_model = hidden_states.shape[-1]
n_layers_plus1 = hidden_states.shape[1]
n_layers = n_layers_plus1 - 1
if layer_range is None:
layer_range = (0.2, 0.8)
start_idx = max(1, int(layer_range[0] * n_layers))
end_idx = min(n_layers, int(layer_range[1] * n_layers))
# First compute regular churn PCs
layer_diffs: Float[Tensor, "batch n_layers d"] = hidden_states[:, 1:, :] - hidden_states[:, :-1, :]
layer_diffs_flat: Float[Tensor, "n d"] = layer_diffs.reshape(-1, d_model).float()
layer_diffs_centered = layer_diffs_flat - layer_diffs_flat.mean(dim=0, keepdim=True)
_, S, Vh = torch.linalg.svd(layer_diffs_centered, full_matrices=False)
n_candidates = min(top_k * 3, Vh.shape[0])
V_candidates: Float[Tensor, "d k"] = Vh[:n_candidates, :].T
proj_mag_sq = (hidden_states.float() @ V_candidates) ** 2
early_mag = proj_mag_sq[:, start_idx:start_idx+3, :].mean(dim=(0, 1))
late_mag = proj_mag_sq[:, end_idx-3:end_idx, :].mean(dim=(0, 1))
mag_slope = late_mag - early_mag # negative = suppressive
# Select top-k by suppressiveness (negative slope)
suppressive_mask = mag_slope < 0
if suppressive_mask.sum() < top_k:
logger.warning(f"Only {suppressive_mask.sum()} suppressive PCs found, taking {top_k} most suppressive")
sorted_indices = torch.argsort(mag_slope, descending=False)[:top_k] # Most negative first
else:
suppressive_indices = torch.where(suppressive_mask)[0]
scores = -mag_slope[suppressive_indices] * S[suppressive_indices] # Higher = more suppressive
sorted_by_score = torch.argsort(scores, descending=True)[:top_k]
sorted_indices = suppressive_indices[sorted_by_score]
V_suppressive: Float[Tensor, "d k"] = V_candidates[:, sorted_indices].to(dtype).to(device).detach()
n_negative = (mag_slope[sorted_indices] < 0).sum().item()
logger.info(f"Churn_suppressive subspace: rank={V_suppressive.shape[1]}, {n_negative}/{top_k} strictly suppressive")
return Subspace(V_suppressive, name="churn_suppressive")
def compute_task_diff_constructive_from_hidden_states(
hidden_states: Float[Tensor, "batch n_layers_plus1 d_model"],
top_k: int =256,
layer_range: Optional[tuple] = None,
) -> Subspace:
"""Compute constructive task_diff: task-discriminative directions being AMPLIFIED.
Standard task_diff is unsigned PCA of (h_cho - h_rej). This variant filters to
directions where the cho/rej separation is INCREASING across layers - i.e., the
model is actively building this distinction, not inheriting it from embeddings.
Method: For each task_diff PC, compute slope of |h_cho @ v| - |h_rej @ v| across layers.
Keep only PCs where separation is growing (constructive discrimination).
Args:
hidden_states: [batch, n_layers+1, d_model] with interleaved cho/rej pairs
top_k: Number of components
layer_range: Optional (start_frac, end_frac) for slope (default 0.3-0.8)
Returns:
Subspace of constructively-discriminating task directions
"""
device = hidden_states.device
dtype = hidden_states.dtype
d_model = hidden_states.shape[-1]
n_layers_plus1 = hidden_states.shape[1]
n_layers = n_layers_plus1 - 1
if layer_range is None:
layer_range = (0.3, 0.8)
start_idx = max(1, int(layer_range[0] * n_layers))
end_idx = min(n_layers, int(layer_range[1] * n_layers))
# Extract cho and rej (interleaved)
hs_cho: Float[Tensor, "n_pairs layers d"] = hidden_states[::2]
hs_rej: Float[Tensor, "n_pairs layers d"] = hidden_states[1::2]
# First compute regular task_diff PCs (on mean diff across layers)
task_diffs: Float[Tensor, "n_pairs d"] = (
hs_cho[:, start_idx:end_idx+1, :] - hs_rej[:, start_idx:end_idx+1, :]
).mean(dim=1).float()
# Per-sample normalize: each pair votes equally regardless of cho-rej magnitude.
# Without this, pairs with large ||cho - rej|| dominate PCA.
task_diffs_norm = normalize_rows(task_diffs)
task_diffs_centered = task_diffs_norm - task_diffs_norm.mean(dim=0, keepdim=True)
_, S, Vh = torch.linalg.svd(task_diffs_centered, full_matrices=False)
n_candidates = min(top_k * 3, Vh.shape[0])
V_candidates: Float[Tensor, "d k"] = Vh[:n_candidates, :].T # [d_model, n_candidates]
# For each PC, compute magnitude separation trend across layers
# |h_cho @ v| - |h_rej @ v| should increase for constructive directions
proj_cho = (hs_cho.float() @ V_candidates).abs() # [n_pairs, n_layers+1, n_candidates]
proj_rej = (hs_rej.float() @ V_candidates).abs()
separation = proj_cho - proj_rej # positive = cho more aligned
# Compute slope: early vs late separation
early_sep = separation[:, start_idx:start_idx+3, :].mean(dim=(0, 1)) # [n_candidates]
late_sep = separation[:, end_idx-3:end_idx, :].mean(dim=(0, 1))
sep_slope = late_sep - early_sep # positive = constructive (separation growing)
# Also check that the direction is actually discriminative (|late_sep| > threshold)
discriminative = late_sep.abs() > 0.01 # Nonzero separation
# Select: constructive AND discriminative
valid_mask = (sep_slope > 0) & discriminative
if valid_mask.sum() < top_k:
logger.warning(f"Only {valid_mask.sum()} constructive+discriminative PCs, taking {top_k} best")
scores = sep_slope * late_sep.abs() # Favor growing + large separation
sorted_indices = torch.argsort(scores, descending=True)[:top_k]
else:
valid_indices = torch.where(valid_mask)[0]
scores = sep_slope[valid_indices] * S[valid_indices]
sorted_by_score = torch.argsort(scores, descending=True)[:top_k]
sorted_indices = valid_indices[sorted_by_score]
V_constructive: Float[Tensor, "d k"] = V_candidates[:, sorted_indices].to(dtype).to(device).detach()
n_valid = ((sep_slope[sorted_indices] > 0) & (late_sep[sorted_indices].abs() > 0.01)).sum().item()
logger.info(f"Task_diff_constructive subspace: rank={V_constructive.shape[1]}, {n_valid}/{top_k} constructive+discriminative")
return Subspace(V_constructive, name="taskdiff_constructive")
def compute_task_diff_from_hidden_states(
hidden_states: Float[Tensor, "batch n_layers_plus1 d_model"],
top_k: int =256,
@@ -1373,88 +795,6 @@ def compute_task_diff_from_hidden_states(
return Subspace(sub.V, name="taskdiff", S=sub.S)
def compute_task_read_subspace(
task_diff_subspace: Subspace,
read_subspace: Subspace,
top_k: int = 256,
) -> Subspace:
"""Compute task_read subspace: task signal readable by transformer blocks.
task_read = task_diff ∩ read
These are task-discriminative directions that are read by attention/MLP inputs
(q/k/v projections, up/gate projections, etc.).
Args:
task_diff_subspace: Subspace of task differences
read_subspace: Subspace readable by residual readers
top_k: Number of components
Returns:
Subspace of task signal that read-modules can read
"""
task_read = approx_intersection(task_diff_subspace, read_subspace)
if task_read.rank > top_k:
task_read = Subspace(task_read.V[:, :top_k], name="taskdiff_read")
else:
task_read.name = "taskdiff_read"
return task_read
def compute_task_lm_head_subspace(
task_diff_subspace: Subspace,
lm_head_subspace: Subspace,
top_k: int = 256,
) -> Subspace:
"""Compute taskdiff_x_logits_read subspace: task signal readable by lm_head.
taskdiff_x_logits_read = taskdiff ∩ logits_read
These are task-discriminative directions that lm_head can read,
i.e. they affect output logits.
"""
taskdiff_logits_read = approx_intersection(task_diff_subspace, lm_head_subspace)
if taskdiff_logits_read.rank > top_k:
taskdiff_logits_read = Subspace(taskdiff_logits_read.V[:, :top_k], name="taskdiff_x_logits_read")
else:
taskdiff_logits_read.name = "taskdiff_x_logits_read"
return taskdiff_logits_read
def compute_task_wnr_subspace(
task_diff_subspace: Subspace,
write_not_read_subspace: Subspace,
top_k: int =256,
) -> Subspace:
"""Compute task_wnr subspace: task signal written but not read.
task_wnr = task_diff ∩ write_not_read
These are task-discriminative directions that are written to residual
but not read by later layers or lm_head.
Args:
task_diff_subspace: Subspace of task differences
write_not_read_subspace: Subspace of write-not-read directions
top_k: Number of components
Returns:
Subspace of task signal that's written but ignored
"""
taskdiff_write_not_read = approx_intersection(task_diff_subspace, write_not_read_subspace)
if taskdiff_write_not_read.rank > top_k:
taskdiff_write_not_read = Subspace(taskdiff_write_not_read.V[:, :top_k], name="taskdiff_x_write_not_read")
else:
taskdiff_write_not_read.name = "taskdiff_x_write_not_read"
return taskdiff_write_not_read
def compute_module_subspace_from_svds(
layer_svds: Dict[str, tuple],
layer_info: Dict[str, dict],