An introduction to graphical tensor notation for mechanistic interpretability

I. TENSORS

• Each tensor leg represents an index, and specifying values for its legs selects an entry in the tensor.For example, A[0,2,0] selects the entry at indices (0, 2, 0).Source: §I · p. 3
• Because storage grows exponentially with the number of legs, high-order tensors are usually represented implicitly as operations among smaller tensors.A tensor with N legs of dimension d contains d^N numbers.Source: §I · p. 3

II. OPERATIONS

• Graphical notation represents tensor products by placing tensors nearby and contractions by connecting legs that share an index and are summed over.The number of free legs indicates the number of legs in the resulting tensor.Source: §II · p. 4
• Rearranging, grouping, or splitting legs depicts index permutations and tensor reshapes, while contractions can be performed in different orders with different efficiencies.For large networks, diagrams make contracted pairs and the result’s structure easier to see than a long index expression.Source: §II · pp. 5–6
• For independent tensors in a tensor network, differentiating the result with respect to one tensor amounts to removing it from the network.The passage attributes this to the multilinearity of einsums.Source: §II · p. 6

III. SPECIAL KINDS OF TENSORS

• An identity matrix is drawn as a bare line, and a delta tensor generalizes this convention by being nonzero only when its indices match.The three-leg delta tensor has ones only along the i = j = k diagonal.Source: §III · pp. 6–7
• Contraction order can change cost dramatically: an order with at most three-leg intermediates scales linearly in tensor count, while finding the optimal order is NP-hard and contraction is #P-hard in general.Simple contraction-order heuristics and approximation techniques can still be useful.Source: §III · p. 6
• Triangles denote isometric maps that preserve vector lengths, with the tip pointing toward the smaller dimension; square isometries are orthogonal or unitary matrices.The reverse contraction does not generally give an identity for a nonsquare matrix because mapping from high to low dimensions squashes some vectors.Source: §III · p. 7
• A Kronecker product of identity matrices is itself an identity matrix after flattening the paired indices.The example gives I_5 ⊗ I_3 = I_15.Source: §III · p. 7

IV. DECOMPOSITIONS (SVD, CP, TUCKER)

• SVD expresses a matrix as two isometric matrices and a diagonal matrix, or equivalently as a sum of outer products weighted by singular values.Small singular values can be discarded to lower the matrix rank; the figure illustrates image approximations retaining 1, 7, 30, or 100 of 200 singular values.Source: §IV · p. 8, Figure 2 · p. 9
• Keeping the largest k singular values gives the best possible rank-k approximation of a matrix under any unitarily invariant norm.The approximation error is determined by the total weight of the discarded singular values.Source: §IV · p. 8
• CP extends the SVD pattern to tensors with more legs, while Tucker relaxes CP by allowing a core tensor that is not restricted to be diagonal.These decompositions work directly on higher-order tensors without first grouping legs into an effective matrix.Source: §IV · pp. 8–9
• Unlike SVD, CP and Tucker decompositions are less well behaved: determining CP-rank is NP-hard, rank is usually chosen manually, and computation usually requires iterative optimization.The passage also notes that isometric-matrix constraints are usually relaxed in these decompositions.Source: §IV · p. 9

V. TENSOR NETWORK DECOMPOSITIONS

• A tensor train represents an exponentially large-dimensional tensor as a chain of smaller tensors, and can be constructed by a sequence of SVDs.Each tensor has a free physical leg and bond legs connected to neighboring tensors.Source: §V · p. 10
• Making tensors isometric toward an orthogonality center makes the network equivalent to an SVD around that center, enabling compression by discarding small singular values and fixing gauge freedom.The stated benefits include sometimes interpretable dominant singular vectors.Source: §V · pp. 10–11
• Low-rank tensor-train compression requires only a few dominant singular vectors at each SVD; otherwise, bond dimensions grow exponentially away from the edges.These decompositions work best when correlations between distant sites that cannot be explained by nearby sites are relatively weak.Source: §V · p. 11
• Tensor networks have gauge freedom because different tensor networks can contract to the same tensor, and inserting an identity and its inverse on a bond can change the network without changing its represented tensor.The passage notes that this can even increase bond dimension arbitrarily when the matrices are rectangular.Source: §V · p. 11

• Neural networks require extending standard tensor-network notation to represent nonlinearities, because that notation disallows nonlinearities and even copying.Source: §VI · p. 12

A. Dense neural networks

• Without nonlinearities, weight matrices can be contracted in any order, or combined into one matrix independently of the input vector.Source: A. Dense neural networks · p. 12
• Nonlinearities are shown as bubbles that require everything inside to be contracted before the function is applied, fixing the contraction order.Source: A. Dense neural networks · p. 12
• Bias terms can be incorporated by concatenating to the input vector and expanding the weight matrices.Source: A. Dense neural networks · p. 12

VII. TRANSFORMERS

• GPT-2’s blocks alternate attention and MLP operations around a residual stream, the network’s main communication channel, into which layers add or subtract transformed information.Source: §VII · p. 13, A. Input and output · p. 16
• The diagram contracts from input at the top to output at the bottom and depicts one of 12 layers; the others repeat its blocks with different learned weights.Source: Figure 4 · p. 14

A. Input and output

• The output gives log probabilities for possible next tokens at every position, with the final position representing the distribution for the unknown next word.Source: A. Input and output · p. 15
• The embedding layer projects each token from vocabulary space into the hidden-dimensional residual stream, while positional embeddings add a position-specific vector.Source: A. Input and output · pp. 15–16
• For simplicity, the document contracts token and positional embeddings into one tensor, though the expanded form can be substituted.Source: A. Input and output · p. 16

B. MLP

• MLP layers contract only the hidden index, acting independently on each token; their up- and down-projection matrices contain most transformer parameters and seem to store many trained facts.The text connects this interpretation to work on key-value memories [31], ROME [32], and activation patching.Source: B. MLP · pp. 16–17

C. Attention

• Attention heads contribute independently to the residual stream, and their sum can be represented in several equivalent ways.Source: C. Attention · pp. 17, 19
• Freezing the attention pattern makes the attention block completely linear, with copying and adding back to the residual stream but no products between copies.Source: C. Attention · p. 18
• Attention patterns are the only network transformation acting on sequence positions, so they determine how information moves between tokens.Source: C. Attention · pp. 18–19
• The low-rank-operator view of attention is probably less useful than representing attention as a sum of heads.Source: C. Attention · p. 19

VIII. COMPOSITION AND PATH EXPANSION

• The analysis focuses on attention-head composition across layers, ignoring MLP blocks and small nonlinearities such as layer normalization.The layers share a residual stream, allowing later layers to use information computed earlier.Source: §VIII · p. 20
• The three simplest nontrivial attention compositions are Q-, K-, and V-composition, although nonlinearities make the expansion approximate.The attention nonlinearities include masking, rescaling, and softmax; layer normalization is omitted from the diagrams and can be treated as approximately linear with care.Source: §VIII · p. 21
• V-composition has a simpler iteration structure than Q- or K-composition: attention patterns and weight transformations are matrix-multiplied.The passage describes the attention patterns and the WO and WV transformations as matrix multiplications.Source: §VIII · p. 22
• A two-layer expansion has ten final-answer terms, and the number of terms grows exponentially with network depth.The ten comprise three non-compositions, three single compositions, and four higher-order compositions.Source: §VIII · p. 22
• Small head dimensions motivate expecting some heads to act mostly independently, but whether low-order terms explain relevant behavior requires empirical checking.The text also notes that other decompositions might explain more behavior with fewer interpretable terms, including Automated Circuit DisCovery (ACDC) [34].Source: §VIII · pp. 22–23

IX. EXAMPLE: TOY INDUCTION HEAD

• The toy circuit uses K-composition of heads in different layers to detect what should follow a repeated token pattern.For example, the earlier “fox” should be retrieved at the final “brown” in “The quick brown fox [...] the quick brown”.Source: §IX · p. 23
• The matching component can use WO and WV in layer 1 and WQ and WK in layer 2 as identities in a relevant subspace, producing a semantic “Match” tensor.“Match” should be near zero when contracted with two vectors that do not semantically match.Source: §IX · p. 24
• The toy construction ignores positional embeddings and may also need to subtract the original final “brown” vector from the residual stream.These are explicitly noted as omissions in the example.Source: §IX · p. 24
• A fixed off-diagonal delta attention pattern maps a matched position to the following token; equivalently, it attends to the previous token.This “previous token” pattern also removes the unwanted final-position term because the attention matrix’s final column is zero.Source: §IX · p. 25
• The retrieved “fox” reaches the residual stream through WV and WO, and WU maps that information to the predicted token.The desired output depends on these matrices handling the information appropriately.Source: §IX · p. 26
• The resulting virtual induction head can have a simple approximate form composed almost entirely of delta-like tensors.The section presents this as a possible form, not a claim that real induction heads are exact copies.Source: §IX · p. 26
• A numerical check on repeated random-vector sequences produces an attention pattern that looks like a real induction-head pattern.The section describes this as a sanity check of the toy construction.Source: §IX · p. 26
• K-composition is only one path-expansion term, so the weight matrices must suppress unwanted terms for the virtual head to become important.WQ, WK, WO, WV, and WU selectively determine what is taken from and added to the residual stream.Source: §IX · p. 27

X. CONCLUSION

• Graphical tensor notation helps communicate interpretability results involving multiple tensors and makes dualities and equivalences easier to see without necessarily losing mathematical rigor.The author presents its usefulness as a preference and argues it remains intuitive.Source: §X · p. 28