If we think parts of MLP0-MLP3 are computing [a sorting algorithm], we can replace those parts with [a sorting algorithm] and check reconstruction loss.[1] However, reconstruction loss is not enough.
Suppose we replace MLP0 with two things:
Its mean activation - simple, but poor reconstruction
MLP0 - perfect reconstruction, but no reduction in complexity
We can visualize this as a pareto frontier trading off reconstruction with "simplicity". Ideally we achieve perfect reconstruction with perfect simplicity.[2] For more intuition on the pareto frontier, we could have an MLP that clusters all inputs in two clusters: "early positions" and "late positions", which would be slightly more complex than the mean.[3]
We can make this an RLVR environment, if only we could clearly...
Define "Simplicity"
Defining simplicity has been complex. But what do we want from a perfectly decomposed model? If we've "perfectly decomposed" a model, then I'd expect ideal circuits to fall out, with "ideal" meaning:
Help predict OOD behavior
Given an [addition] circuit, we can know which types of inputs it'll succeed & fail on (and why)
Be extractable & minimal
The smallest part of the model that does [addition]
Be removable w/ minimal harm to unrelated circuits
Affects [addition] but not unrelated tasks like [bracket closing]
Currently, I believe we want "circuit simplicity" defined as having a small number of nodes and edges:
Nodes - variables like "numbers", "dog-like", "within quotation marks"
Edges - causal links. Ideally "dog-like" is computed by a small number of earlier features (eg "dog ear", "dog tail", etc).
Additionally, an ideal decomposition should tell us which features are reused. If two nodes affect all the same downstream nodes, then you can combine them. This reduces both the number of nodes and edges (similar to refactoring code).
In contrast, a complex, dense graph would be two MLPs.
MLP = D(ReLU(Wx))
We can define the "nodes/features" as the hidden dimension:
Nodes = ReLU(Wx)
These nodes are bad because they're polysemantic: they have different meanings at different activation ranges.
When you look at which Nodes in layer 1 are causally important for Nodes in layer 2, it's a dense graph (ie each node is caused by a large number of previous nodes).
Tensor Networks Don't Solve This Issue
In the simpler case of two bilinear layers, you can simply matrix multiply them together and see exactly which hidden dimensions/nodes of the first cause the second directly from the weights. Sparse connections can even be directly optimized for this way.
HOWEVER, we still have a dense graph. The main problem here is our model isn't basis-aligned (ie each coordinate of activations being monosemantic).
With that said, I believe that tensor networks give us better tools to find the correct basis, I am simply too confused atm to make the most of it.[4]
Red Herrings of Simplicity
Bits of information - You can quantize a model into fewer bits of information (float32 to float16), but I won't understand the model better than I used to (so is a bad metric as stated).[5]
Alternatively, I do think an MLP mapping inputs into two clusters is indeed simpler than one that maps to 50k clusters which feels like bits-of-info related.
Low Rank - If my mlp only writes to/reads from a rank 20 subspace, that seems useful to know (because it constrains what other MLPs/attn heads it interacts with) but it's not simpler.
eg If I know a part of the model does [parity], then it's not important if it's rank 2 or 20. If they're caused by the same number of upstream components (and vice versa for downstream components), then that's same number of edges.
Low Features per Datapoint - SAEs & PD (parameter decomposition) care about this. PD cares about the smallest number of "mechanisms" for the full computation path input-to-output.
On its own it fails (define the whole network as 1 feature); you also want a small number of total features to choose from. In hand-wavy terms, we want a simple dictionary (choice of features) and for datapoints to be simply explained, conditioning on that dictionary (...I think?).
Functional Complexity (?) - I have bilinear layers that are a sum of 1k quadratic equations. Maybe you only need 200 quadratic functions and the rest can be approximated as linear (I think this is helpful, but am unsure).
How to Gain Tractability
If we know how models represent [redundancy, circuit suppression, compositionality, hierarchical structure, DAGs, etc], we can decompose modules by just looking at the weights (using tensor transformers). But we don't know those things, so we don't know the right ontology to use when defining "simplicity". There are ways we can gain traction though!
Tract 1: Death Success by 1000 Circuits
With many circuits in a real LLM, we can find components that are shared across circuits. They might read from the same subspace (activation-based) or use the same MLP (weight-based). However, whole MLPs (and whole attn-heads) are the wrong fundamental units to use. Two circuits may use the same MLP but different parts of it.
If we understood which parts of an MLP 1000 circuits shared and didn't share, then that'd be our decomposition. We would then have a ground truth label for how we should've decomposed that MLP eg "Oh there's clearly some hierarchical structure here once you see how it interacts with later attn-heads". Then we can directly decompose the model from the weights from the beggining, with all the circuits falling out naturally.
Tract 2: QK OV Circuits but for Everything
In Anthropic's Mathematical Framework for Transformer Circuits, they introduced QK & OV circuits. They stopped there because you can't compose other parts of the model together due to nonlinearities such as softmax or GeLU. With a tensor-transformer, you can compose everything with everything, such as embedding --> attn_0 --> mlp_0 --> unembedding.
We can compose any intermediate attn/MLP components together, for any computational path through the network, forming a new sub-tensor that is that computational path. That sub-tensor might have structure (hierarchical, DAG, etc) that individual components didn't.
This property of compositionality is extremely powerful!
Another way to take advantage of this is to take a specific vector from:
An unembedding row (eg "myself")
An unembedding diff (eg "myself" - "yourself")
Classes in the unembedding (eg numbers "1", "2", ... might be in a similar subspace)
Steering vector
You take that vector/subspace and make that the end point of your computational path (ie you can matrix-multiply that vector with any intermediate matrix from Attn_out or MLP_out). Tensor transformers allow you to specify precisely any hypothesis, BUT they don't tell you which hypotheses you should consider, hence I haven't solved mech interp yet.
Tract 3: Interpreting Small Models
We can have small models trained on algorithmic tasks (algzoo) which already exhibit many behaviors we don't know how to approach. For example, training on CE loss on softmax-ed logits incentivizes circuit suppression: bird features can upweight the "bird" class AND/OR downweight other classes. It's more about a class's relative logit value compared to other classes than its objective value. AFAIK, we don't have a solid understanding of how models represent this.
Additionally, models are likely a mix of compositional functions/features and memorization/look-up tables. We can directly tackle this like Linda Linsefors does by training a toy model of memorization. This allows us to better analyze the model in silico as well as build out tools to detect and tease out the structure, while varying hyperparameters.
If we can deconfuse ourselves enough, we can fully realize mech interp's potential as a verifiable task. This would be a big deal!
We can train models of varying layers, model width, datasets, taking checkpoints and THEN get SOTA models to solve this with RLVR. With enough data of models being reverse-engineered, we can solve training dynamics, how data leads to algorithms, task-ai, etc.
We'll still need to solve value-alignment, of course, but maybe that could be a tensor network too?[6]
If you would like to help in this project, you can always dm me here or on discord at #loganriggs. I'm currently working out of this repo trying to reverse engineer the 500M parameter tensor transformer I trained last year.
The mean or a constant bias are both quite interpretable. They're the same for all inputs, so we don't have to reason about "conditioning on these tokens previously, the output will change like so". It's a constant bias, dead simple.
Which decomposition method we select is itself the work of much of Mech Interp, no? And the evaluation of a decomposition method is based on much fuzzier criteria.
If we think parts of MLP0-MLP3 are computing [a sorting algorithm], we can replace those parts with [a sorting algorithm] and check reconstruction loss.[1] However, reconstruction loss is not enough.
Suppose we replace MLP0 with two things:
We can visualize this as a pareto frontier trading off reconstruction with "simplicity". Ideally we achieve perfect reconstruction with perfect simplicity.[2] For more intuition on the pareto frontier, we could have an MLP that clusters all inputs in two clusters: "early positions" and "late positions", which would be slightly more complex than the mean.[3]
We can make this an RLVR environment, if only we could clearly...
Define "Simplicity"
Defining simplicity has been complex. But what do we want from a perfectly decomposed model? If we've "perfectly decomposed" a model, then I'd expect ideal circuits to fall out, with "ideal" meaning:
Currently, I believe we want "circuit simplicity" defined as having a small number of nodes and edges:
Additionally, an ideal decomposition should tell us which features are reused. If two nodes affect all the same downstream nodes, then you can combine them. This reduces both the number of nodes and edges (similar to refactoring code).
In contrast, a complex, dense graph would be two MLPs.
MLP = D(ReLU(Wx))
We can define the "nodes/features" as the hidden dimension:
Nodes = ReLU(Wx)
These nodes are bad because they're polysemantic: they have different meanings at different activation ranges.
When you look at which Nodes in layer 1 are causally important for Nodes in layer 2, it's a dense graph (ie each node is caused by a large number of previous nodes).
Tensor Networks Don't Solve This Issue
In the simpler case of two bilinear layers, you can simply matrix multiply them together and see exactly which hidden dimensions/nodes of the first cause the second directly from the weights. Sparse connections can even be directly optimized for this way.
HOWEVER, we still have a dense graph. The main problem here is our model isn't basis-aligned (ie each coordinate of activations being monosemantic).
With that said, I believe that tensor networks give us better tools to find the correct basis, I am simply too confused atm to make the most of it.[4]
Red Herrings of Simplicity
How to Gain Tractability
If we know how models represent [redundancy, circuit suppression, compositionality, hierarchical structure, DAGs, etc], we can decompose modules by just looking at the weights (using tensor transformers). But we don't know those things, so we don't know the right ontology to use when defining "simplicity". There are ways we can gain traction though!
Tract 1:
DeathSuccess by 1000 CircuitsWith many circuits in a real LLM, we can find components that are shared across circuits. They might read from the same subspace (activation-based) or use the same MLP (weight-based). However, whole MLPs (and whole attn-heads) are the wrong fundamental units to use. Two circuits may use the same MLP but different parts of it.
If we understood which parts of an MLP 1000 circuits shared and didn't share, then that'd be our decomposition. We would then have a ground truth label for how we should've decomposed that MLP eg "Oh there's clearly some hierarchical structure here once you see how it interacts with later attn-heads". Then we can directly decompose the model from the weights from the beggining, with all the circuits falling out naturally.
Tract 2: QK OV Circuits but for Everything
In Anthropic's Mathematical Framework for Transformer Circuits, they introduced QK & OV circuits. They stopped there because you can't compose other parts of the model together due to nonlinearities such as softmax or GeLU. With a tensor-transformer, you can compose everything with everything, such as embedding --> attn_0 --> mlp_0 --> unembedding.
We can compose any intermediate attn/MLP components together, for any computational path through the network, forming a new sub-tensor that is that computational path. That sub-tensor might have structure (hierarchical, DAG, etc) that individual components didn't.
This property of compositionality is extremely powerful!
Another way to take advantage of this is to take a specific vector from:
You take that vector/subspace and make that the end point of your computational path (ie you can matrix-multiply that vector with any intermediate matrix from Attn_out or MLP_out). Tensor transformers allow you to specify precisely any hypothesis, BUT they don't tell you which hypotheses you should consider, hence I haven't solved mech interp yet.
Tract 3: Interpreting Small Models
We can have small models trained on algorithmic tasks (algzoo) which already exhibit many behaviors we don't know how to approach. For example, training on CE loss on softmax-ed logits incentivizes circuit suppression: bird features can upweight the "bird" class AND/OR downweight other classes. It's more about a class's relative logit value compared to other classes than its objective value. AFAIK, we don't have a solid understanding of how models represent this.
Additionally, models are likely a mix of compositional functions/features and memorization/look-up tables. We can directly tackle this like Linda Linsefors does by training a toy model of memorization. This allows us to better analyze the model in silico as well as build out tools to detect and tease out the structure, while varying hyperparameters.
Again shoehorning tensor networks, if we can state our hypothesis as a tensor, we can directly compare from the model weights whether they're equivalent.
Big if True
If we can deconfuse ourselves enough, we can fully realize mech interp's potential as a verifiable task. This would be a big deal!
We can train models of varying layers, model width, datasets, taking checkpoints and THEN get SOTA models to solve this with RLVR. With enough data of models being reverse-engineered, we can solve training dynamics, how data leads to algorithms, task-ai, etc.
We'll still need to solve value-alignment, of course, but maybe that could be a tensor network too?[6]
If you would like to help in this project, you can always dm me here or on discord at #loganriggs. I'm currently working out of this repo trying to reverse engineer the 500M parameter tensor transformer I trained last year.
Explaining Attention with Program Synthesis does this with program synthesis, replacing attention patterns with 1000 programs for up to a 3B param model.
This is the case for one Attn head in my model that is basically a bias (for outlier dimension reasons).
The mean or a constant bias are both quite interpretable. They're the same for all inputs, so we don't have to reason about "conditioning on these tokens previously, the output will change like so". It's a constant bias, dead simple.
This post is partially a cry for help; I've only worked on this for 2 months, but would appreciate more eyes on this problem.
Plausibly having only -1, 0, & 1 be valid weights might make our job easier though. Haven't thought too much about this.
jk jk. But maybe??