TransformerLens + PyTorch Quick Reference

ARENA Mech Interp Week - J Rosser

1 TransformerLens: Model Loading 3.2 Forward with Cache 5 TransformerLens: Weight Access
1.1 Load Pretrained Model logits , cache = model . run _with_c ache ( tokens ) # Embeddings
model . W_E # [ d_vocab , d_model ]
# Cache only specific activations model . W_pos # [ n_ctx , d_model ]
from t r a n s f o r m e r _ l e n s import H o o k e d T r a n s f o r m e r logits , cache = model . run _with_c ache ( model . W_U # [ d_model , d_vocab ]
tokens ,
model = H o o k e d T r a n s f o r m e r . f r om _p re t ra in e d ( names_filter = lambda n : n . endswith ( " pattern " ) ,
" gpt2 - small " , # Attention weights ( all have layer , head dims )
return_type = None # Don ’t return logits model . W_Q [ layer , head ] # [ d_model , d_head ]
cente r_unemb ed = True , # Center W_U )
c e n t e r _ w r i t i n g _ w e i g h t s = True , model . W_K [ layer , head ] # [ d_model , d_head ]
fold_ln = True , # Fold LayerNorm model . W_V [ layer , head ] # [ d_model , d_head ]
r e f a c t o r _ f a c t o r e d _ a t t n _ m a t r i c e s = True , 3.3 Forward with Hooks model . W_O [ layer , head ] # [ d_head , d_model ]
device = device
) def my_hook ( activation , hook ) : # MLP weights
# Modify activation model . W_in [ layer ] # [ d_model , d_mlp ]
model . W_out [ layer ] # [ d_mlp , d_model ]
1.2 Create Custom Model return activation # Must return if modifying
model . b_in [ layer ] # [ d_mlp ]
logits = model . run_w ith_hoo ks ( model . b_out [ layer ] # [ d_model ]
from t r a n s f o r m e r _ l e n s import H o o k e d T r a n s f o r m e r C o n f i g tokens ,
fwd_hooks =[ # Composed circuits
cfg = H o o k e d T r a n s f o r m e r C o n f i g ( ( utils . get_act_name ( " z " , 0) , my_hook ) , W_OV = model . W_V [l , h ] @ model . W_O [l , h ] # OV circuit
n_layers =2 , n_heads =12 , d_model =768 , ( " blocks .1. attn . hook_pattern " , other_hook ) , W_QK = model . W_Q [l , h ] @ model . W_K [l , h ]. T # QK circuit
d_head =64 , d_mlp =3072 , d_vocab =50257 , ]
n_ctx =1024 , act_fn = " gelu " , )
n o r m a l i z a t i o n _ t y p e = " LN " , # or " LNPre " , None
attention_dir = " causal " , # or " bidirectional "
6 TransformerLens: Hook Names
)
device = device
4 TransformerLens: Cache Access from transformer_lens import utils
model = H o o k e d T r a n s f o r m e r ( cfg )
# Get hook name for activation type
4.1 Access Patterns utils . get_act_name ( " pattern " , layer )
1.3 Model Config Access # Tuple syntax ( preferred )
utils . get_act_name ( " z " , layer )
utils . get_act_name ( " result " , layer )
model . cfg . n_layers # Number of layers cache [ " pattern " , layer ] # Attention patterns utils . get_act_name ( " resid_pre " , layer )
model . cfg . n_heads # Heads per layer cache [ " q " , layer ] # Query vectors utils . get_act_name ( " resid_post " , layer )
model . cfg . d_model # Model dimension cache [ " k " , layer ] # Key vectors utils . get_act_name ( " q " , layer )
model . cfg . d_head # Head dimension cache [ " v " , layer ] # Value vectors utils . get_act_name ( " k " , layer )
model . cfg . d_mlp # MLP dimension cache [ " z " , layer ] # Pre - projection output utils . get_act_name ( " v " , layer )
model . cfg . d_vocab # Vocabulary size cache [ " result " , layer ] # Post - projection output utils . get_act_name ( " mlp_out " , layer )
model . cfg . n_ctx # Context length cache [ " resid_pre " , layer ] # Residual before layer utils . get_act_name ( " post " , layer ) # MLP post - act
model . cfg . device # Device cache [ " resid_mid " , layer ] # After attn , before MLP
cache [ " resid_post " , layer ] # After layer # Common hook name patterns
cache [ " attn_out " , layer ] # Attention output " hook_embed "
2 TransformerLens: Tokenization cache [ " mlp_out " , layer ]
cache [ " post " , layer ]
#
#
MLP output
MLP post - activation
" hook_pos_embed "
" blocks .{ L }. hook_resid_pre "
cache [ " pre " , layer ] # MLP pre - activation " blocks .{ L }. attn . hook_pattern "
# Text -> Token IDs ( prepends BOS by default ) " blocks .{ L }. attn . hook_z "
tokens = model . to_tokens ( " Hello world " ) # String syntax " blocks .{ L }. attn . hook_result "
tokens = model . to_tokens ( text , prepend_bos = False ) cache [ " blocks .0. attn . hook_pattern " ] " blocks .{ L }. hook_mlp_out "
" ln_final . hook_normalized "
# Text -> String tokens list
str_toks = model . to_str_tokens ( " Hello world " ) 4.2 Cache Methods
# [ ’ <| endoftext | > ’ , ’ Hello ’, ’ world ’]
# Apply LayerNorm to stacked residuals 7 TransformerLens: FactoredMatrix
# Token IDs -> String scaled = cache . a p p l y _ l n _ t o _ s t a c k (
text = model . to_string ( tokens ) residual_stack , layer = -1 , pos_slice = -1 from transformer_lens import FactoredMatrix
)
# Single token operations # Create from two matrices ( stores A , B separately )
tok_id = model . t o _s in gl e _t ok en ( " Paris " ) # Get accumulated residual stream OV = FactoredMatrix ( model . W_V [l , h ] , model . W_O [l , h ])
tok_str = model . t o _ s i n g l e _ s t r _ t o k e n ( tok_id ) accum , labels = cache . a c c u m u l a t e d _ r e s i d (
layer = -1 , incl_mid = True , # Chain operations ( stays factored )
# Batch decode pos_slice = -1 , return_labels = True full_OV = model . W_E @ OV @ model . W_U
texts = model . tokenizer . batch_decode ( token_ids ) )
# Materialize only when needed
# Decompose by component full_matrix = OV . AB
3 TransformerLens: Forward Pass decomp , labels = cache . de c om po se _ re si d (
layer = -1 , pos_slice = -1 , return_labels = True # Efficient properties
) OV . eigenvalues # Only non - zero eigenvalues
3.1 Basic Forward Pass # Stack all head outputs
OV . S # Singular values
OV . norm () # Frobenius norm
logits = model ( tokens ) # [ batch , seq , d_vocab ] heads , labels = cache . s t a c k _ h e a d _ r e s u l t s ( OV . shape # Shape of full matrix
loss = model ( tokens , return_type = " loss " ) layer = -1 , pos_slice = -1 , return_labels = True
) # Efficient indexing ( returns FactoredMatrix )
submatrix = full_OV [ indices , indices ]

1
TransformerLens + PyTorch ARENA Quick Reference

8 TransformerLens: Hooks # Move tensors / models tensor . std ( dim )

tensor = tensor . to ( device ) tensor . var ( dim , unbiased = False )
8.1 Hook Function Patterns model = model . to ( device ) tensor . max ( dim )
tensor . min ( dim )
# Returns ( values , indices )

# Disable gradients tensor . argmax ( dim )

from t r a n s f o r m e r _ l e n s . hook_points import HookPoint t . s e t _ g r a d _ e n a b l e d ( False ) tensor . argmin ( dim )
from functools import partial tensor . norm ( dim = dim , keepdim = True )
# Inference mode context
# Access hook ( no return = no modification ) with t . infe rence_m ode () : # keepdim = True preserves dimension
def access_hook ( act , hook : HookPoint ) : logits = model ( tokens )
store [ hook . layer () ] = act . mean () # Softmax family
# No grad context tensor . softmax ( dim = -1)
# Modify hook ( must return tensor ) with t . no_grad () : tensor . log_softmax ( dim = -1)
def modify_hook ( act , hook : HookPoint ) : logits = model ( tokens )
act [: , : , head_idx , :] = 0.0 # Matrix operations
return act tensor @ other # Matrix multiply
# Hook with extra arguments 11 PyTorch: Tensor Creation t . matmul (a , b )
t . bmm (a , b ) # Batched matmul
def param_hook ( act , hook , head_idx , scale ) : t . einsum ( " ij , jk - > ik " , a , b )
act [: , : , head_idx , :] *= scale # From data
return act t . tensor ([1 , 2 , 3]) # Comparison
t . tensor ( data , device = device , dtype = t . float32 ) tensor > 0 # Returns bool tensor
hook_fn = partial ( param_hook , head_idx =4 , scale =0.5) tensor . any ( dim )
# Zeros / Ones tensor . all ( dim )
t . zeros ( batch , seq , d_model )
8.2 Hook Management t . ones ( shape )
t . where ( cond , x , y ) # Conditional select
t . zeros_like ( tensor )
# Reset all hooks
model . reset_hooks ()
model . reset_hooks ( i n c l u d i n g _ p e r m a n e n t = True )
t . ones_like ( tensor )
14 PyTorch: Indexing
# Random
t . rand ( shape ) # Uniform [0 , 1) # Basic slicing
# Permanent hooks t . randn ( shape ) # Normal (0 , 1) tensor [: , -1] # Last position
model . add_hook ( hook_name , hook_fn , is_permanent = True ) t . randint (0 , 100 , ( batch , seq ) ) tensor [: , : -1] # All but last
t . randperm ( n ) # Random permutation tensor [... , idx ] # Ellipsis for batch dims
# Hook context dictionary
hook . ctx [ " key " ] = value # Sequences # Gather ( select along dim )
stored = model . hook_dict [ " hook_name " ]. ctx [ " key " ] t . arange ( start , end , step ) values = tensor . gather ( dim = -1 , index = indices )
t . linspace ( start , end , steps )
# Advanced indexing
9 SAE-Lens Integration # Special
t . empty ( shape ) # Uninitialized
tensor [ bool_mask ]
tensor [ idx_tensor ]
# Boolean indexing
# Integer indexing
from sae_lens import SAE , H o o k e d S A E T r a n s f o r m e r t . full ( shape , value ) tensor [ rows , cols ] # Multi - dim indexing
t . eye ( n ) # Identity matrix
# Load model with SAE support # Topk
model = H o o k e d S A E T r a n s f o r m e r . f ro m_ p re tr ai n ed ( values , indices = tensor . topk (k , dim = -1)

)
" gpt2 - small " , device = device 12 PyTorch: Tensor Operations # Sorting
sorted_vals , indices = tensor . sort ( dim = -1)
# Shape operations indices = tensor . argsort ( descending = True )
# Load pretrained SAE tensor . shape # Get shape
sae , cfg , sparsity = SAE . f r om _p r et ra in e d ( tensor . size ( dim ) # Size of dimension
release = " gpt2 - small - res - jb " , # Unique
tensor . numel () # Total elements unique , counts = t . unique ( tensor , return_counts = True )
sae_id = " blocks .7. hook_ resid_p re " , tensor . squeeze ( dim ) # Remove dim of size 1
device = str ( device ) , tensor . unsqueeze ( dim ) # Add dim of size 1
) # Nonzero
tensor . flatten () # Flatten to 1 D indices = t . nonzero ( tensor > 0)
tensor . flatten ( start , end ) # Flatten range indices = t . argwhere ( tensor > 0)
# Run with SAE ( context manager ) tensor . reshape ( shape )
with model . saes ( saes =[ sae ]) : tensor . view ( shape ) # Must be contiguous
logits = model ( tokens ) # Diagonal
tensor . T # Transpose (2 D ) tensor . diagonal ( offset =0)
tensor . transpose ( d1 , d2 ) tensor . diag ( offset )
# Run with SAE ( explicit ) tensor . permute ( dims )
model . add_sae ( sae )
logits = model ( tokens )
model . reset_saes () # Don ’t forget !
# Stacking / Concatenation
t . stack ([ t1 , t2 ] , dim =0) # New dimension 15 PyTorch: In-Place & Masking
t . cat ([ t1 , t2 ] , dim =0) # Existing dim
# Cache SAE activations t . concat ([ t1 , t2 ]) # Alias for cat # In - place operations ( end with _ )
_ , cache = model . r u n _ w i t h _ c a c h e _ w i t h _ s a e s ( tensor . zero_ ()
tokens , saes =[ sae ] # Splitting tensor . fill_ ( value )
) t . split ( tensor , size , dim ) tensor . masked_fill_ ( mask , value )
sae_acts = cache [ f " { sae . cfg . hook_name }. h o o k _ s a e _ a c t s _ p o s t " ] t . chunk ( tensor , chunks , dim )
tensor . unbind ( dim ) # Returns tuple # Masking patterns
# SAE weights mask = t . triu ( t . ones ( seq , seq ) , diagonal =1) . bool ()
sae . W_enc # [ d_in , d_sae ] encoder tensor . masked_fill_ ( mask , float ( ’ - inf ’) )
sae . W_dec # [ d_sae , d_in ] decoder
13 PyTorch: Math Operations # Clone vs reference
new = tensor . clone () # Copy data
10 PyTorch: Device Management # Elementwise
tensor . abs ()
new = tensor . detach () # Detach from graph
new = tensor . detach () . clone () # Both
tensor . exp ()
import torch as t tensor . log ()
# Device selection
device = t . device (
tensor . sqrt ()
tensor . pow ( n ) 16 PyTorch: Type Conversion
tensor . clip ( min , max )
" mps " if t . backends . mps . is_available () tensor . float () # to float32
else " cuda " if t . cuda . is_available () # Reductions tensor . half () # to float16
else " cpu " tensor . sum ( dim ) tensor . long () # to int64
) tensor . mean ( dim ) tensor . int () # to int32

2
TransformerLens + PyTorch ARENA Quick Reference

tensor . bool () # to bool ): tokens = str_tokens ,

tensor . to ( dtype ) # explicit dtype """ Patch clean activation into corrupted run . """ attention = att ent ion _pa tte rn
corrupted_act [: , pos , :] = clean_cache [ hook . name ][: , pos , :] )
# To Python return corrupted_act
tensor . item () # Scalar -> Python number
tensor . tolist ()
tensor . numpy ()
#
#
To Python list
To numpy ( CPU only )
from functools import partial 23 Jaxtyping Annotations
tensor . cpu () . numpy () # Move to CPU first hook_fn = partial (
patch_hook , pos = end_pos , clean_cache = clean_cache from jaxtyping import Float , Int , Bool
# From numpy ) from torch import Tensor
t . from_numpy ( array ) patc hed_logi ts = model . run_wit h_hooks (
t . tensor ( array ) # Copies data corrupted_tokens , # Common annotations
fwd_hooks =[( utils . get_act_name ( " resid_pre " , layer ) , hook_fn ) ] Float [ Tensor , " batch seq d_model " ]
) Float [ Tensor , " batch heads seq_q seq_k " ]

17 Einops Patterns Int [ Tensor , " batch seq " ]

Bool [ Tensor , " batch " ]

import einops 20 Common Patterns: Ablation # With variable dims

Float [ Tensor , " ... d_model " ]
# Rearrange ( reshape with named dims ) # Zero ablation Float [ Tensor , " * batch seq d_model " ]
einops . rearrange (x , " b s h d -> b s ( h d ) " ) def z e r o _ a b l a t e _ h e a d (z , hook , head_idx ) :
einops . rearrange (x , " b ( h w ) c -> b h w c " , h =28) z [: , : , head_idx , :] = 0.0
einops . rearrange (x , " ( b1 b2 ) ... -> b1 b2 ... " , b1 =2) return z 24 Training Loop Patterns
# Reduce ( aggregation ) # Mean ablation
einops . reduce (x , " b s d -> b d " , " mean " ) def m e a n _ a b l a t e _ h e a d (z , hook , head_idx ) : optimizer = t . optim . AdamW (
einops . reduce (x , " b s h d -> b h " , " sum " ) z [: , : , head_idx , :] = z [: , : , head_idx , :]. mean (0) model . parameters () , lr =1 e -4 , weight_decay =0.01
return z )
# Repeat ( broadcast )
einops . repeat (x , " d -> b d " , b = batch_size ) # Ablate specific heads for epoch in range ( n_epochs ) :
einops . repeat (x , " h d -> h n d " , n = seq_len ) def ablate_heads (z , hook , he ad s_ t o_ ab la t e ) : for batch in dataloader :
for head in he ad s_ t o_ ab la t e : optimizer . zero_grad ()
# Einsum ( tensor contraction ) z [: , : , head , :] = 0.0
einops . einsum (q , k , return z logits = model ( batch )
" b s h d , b t h d -> b h s t " ) loss = compute_loss ( logits , targets )
einops . einsum ( attn , v ,
" b h s t , b t h d -> b s h d " )
einops . einsum ( resid , W_U ,
21 Common Patterns: Direct Logit Attri- loss . backward ()
optimizer . step ()
" b s d , d v -> b s v " )
bution # Optional : learning rate scheduling
for group in optimizer . param_groups :
# Common attention pattern
einops . einsum (q , k , # Get direction in residual stream group [ " lr " ] = new_lr
" batch posQ heads dhead , batch posK heads dhead -> batch heads posQ posK " ) logit_dir = model . W_U [: , correct_tok ] - model . W_U [: , incorrect_tok ]

# Attribute head outputs 25 Useful Imports

18 Common Patterns: Logit Diff head_outputs = cache [ " result " , layer ] # [ batch , seq , heads , d_model ]
dla = einops . einsum ( import torch as t
head_outputs [: , -1] , logit_dir , from torch import Tensor
def l o g i t s _ t o _ a v e _ l o g i t _ d i f f ( " batch heads d_model , d_model -> batch heads "
logits , answer_tokens , per_prompt = False import einops
) from functools import partial
):
""" Compute logit difference metric . """ from tqdm import tqdm
# For SAE latents
final_logits = logits [: , -1 , :] sae_acts = cache [ sae_acts_hook ] # [ batch , seq , d_sae ]
answer_logits = final_logits . gather ( from transformer_lens import (
dla = sae_acts [: , -1] * ( sae . W_dec @ logit_dir ) HookedTransformer ,
dim = -1 , index = answer_tokens
) HookedTransformerConfig ,
ActivationCache ,
correct , incorrect = answer_logits . unbind ( dim = -1)
logit_diff = correct - incorrect 22 CircuitsVis FactoredMatrix ,
utils ,
if per_prompt :
return logit_diff import circuitsvis as cv )
return logit_diff . mean () from transformer_lens . hook_points import HookPoint
# Attention patterns
cv . attention . a t t e n t i o n _ p a t t e r n s ( from sae_lens import SAE , H o o k e d S A E T r a n s f o r m e r

19 Common Patterns: Activation Patch- tokens = model . to_str_tokens ( tokens ) ,

attention = cache [ " pattern " , layer ][0] , from jaxtyping import Float , Int , Bool

ing )
a t t e n t i o n _ h e a d _ n a m e s =[ f " L { l } H { h } " for h in range ( n_heads ) ]
import circuitsvis as cv
import plotly . express as px
def patch_hook ( # Attention heads view
corrupted_act , hook , pos , clean_cache cv . attention . at te nt i on _h e ad s (