From Physics to AI Safety: Cross-Disciplinary Questions for Transformer Interpretability

Transformer Architecture Deep Dive

In this section, we focus on a decoder-only transformer, specifically using the same architecture and parameters as GPT-2 Small [3].

Foundational Concepts

The Residual Stream: Information Highway

At the heart of the transformer lies the residual stream - the main information highway that carries data through each layer. Unlike traditional neural networks where information gets completely transformed at each step, transformers add modifications to this stream while preserving the original information. This is analogous to how we might annotate a document by adding notes in the margins rather than rewriting the entire text. From a physics perspective, this is similar to perturbation theory - we start with an unperturbed state (the original embeddings) and add small corrections at each layer, with each transformer block acting like a perturbative correction that modifies but doesn't replace the underlying state.

Layer Normalization: Preventing Gradient Problems

Layer normalization's primary purpose is maintaining stable gradient flow during training. In deep networks, activations can have wildly different scales across layers, causing gradients during backpropagation to become very small (vanishing) or very large (exploding). Layer normalization solves this by normalizing the activations going into each layer, keeping their scale consistent.

The process: subtract the mean to center at zero, divide by standard deviation to normalize variance, then apply learned scaling and translation parameters. This ensures more stable gradient magnitudes throughout the network during training, allowing the model to learn effectively. The learned scaling (γ) and translation (β) parameters then allow each layer to recover the optimal scale and shift for its specific function, while maintaining the gradient flow benefits.

This creates more stable optimization dynamics, similar to how physicists normalize variables in numerical simulations to prevent computational instabilities when dealing with vastly different scales.

The Transformer Flow

See Fig.2 above starting from bottom to top.

1. Tokens: From Text to Integers

Everything starts with converting human-readable text into integers through tokenization. This uses a method called Byte-Pair Encoding - we start with individual characters and iteratively merge the most common pairs. The result is a vocabulary of ~50,000 sub-word tokens that can represent any text. Each token gets mapped to a unique integer ID, which serves as an efficient key for the embedding lookup table.

2. Embedding: From Integers to Vectors

The embedding layer is essentially a giant lookup table that converts each token ID into a high-dimensional vector (768 dimensions in GPT-2). This learned mapping allows the model to represent semantic relationships - similar tokens end up with similar embedding vectors through training.

3. Positional Embeddings: Adding Sequence Information

Since attention is content-based rather than position-based, we need to explicitly tell the model where each token sits in the sequence. We learn another lookup table that maps position indices to vectors, then add these to the token embeddings. The key insight that initially confused me: how does adding positional embeddings to token embeddings not create collisions? The answer lies in the high dimensionality - with embedding vectors of length 768, the probability of collision becomes vanishingly small. This is similar to how CRISPR achieved precision in gene editing: while earlier techniques used shorter recognition sequences (higher collision risk), CRISPR uses ~20-nucleotide guide sequences, making the probability of accidentally matching the wrong genomic location negligibly small.

4. Attention Mechanism: Intelligent Information Routing

This is where the magic happens. For each token, the model computes three vectors: Query (what am I looking for?), Key (what do I contain?), and Value (what information should I pass along?). The similarity between Query and Key vectors determines how much attention each token pays to every other token in the sequence. (During implementation, I discovered that the order of performing Einstein summation operations within each dot product doesn't affect the results—in retrospect it makes total sense because of commutativity and associativity, but it was a little fun moment.)

**Note: In decoder-only transformers like GPT-2, this is specifically self-attention - each token attends to other tokens within the same sequence. This differs from encoder-decoder architectures (like the original Transformer) which also include cross-attention, where decoder tokens attend to encoder representations. Since we have no separate encoder, all attention is self-attention.**

Multi-Head Attention: Rather than having one attention mechanism, transformers use multiple "heads" (12 in GPT-2) that attend to different types of information simultaneously. Each head can specialize - one might focus on syntactic relationships, another on semantic meaning. It's like having multiple experts each contributing their specialized perspective.

Causal Attention: In GPT-style models, we use what's called "causal attention" - though I find this term somewhat misleading from a physicist's perspective. While there's a notion of sequence, there's no true temporal causality. Instead, it's a masking mechanism that prevents tokens from seeing later tokens in sequence, maintaining the sequential prediction task. This is analogous to preventing data leakage in ML (where future information accidentally gets into training features) or blinding in experimental physics (where researchers are kept unaware of group assignments to prevent bias). All three techniques share the same core principle: controlling information flow to maintain the integrity of the process.

The Complete Attention Process: The mechanism follows six precise steps: (1) compute Query, Key, and Value vectors by applying learned linear transformations to the input embeddings, (2) calculate raw attention scores by computing dot products between each Query and every Key vector, (3) scale these scores to prevent extremely large values, (4) apply the causal mask to set future positions to negative infinity (preventing information leakage), (5) convert masked scores to probabilities using softmax, and (6) compute the weighted sum of Value vectors using these probabilities, then apply a final linear transformation to produce the output that gets added back to the residual stream.

5. MLP Layer: Pattern Recognition and Computation

Despite the confusing name "multi-layer perceptron," these are actually simple two-layer feed-forward networks: expand the dimensionality (768 → 3072), apply non-linearity (GELU activation), then contract back (3072 → 768). The expansion (768 → 3072) creates finer granularity by providing more "working space" for complex transformations. The GELU non-linearity then selectively activates different pathways, creating pattern recognition. The contraction (3072 → 768) distills all this rich intermediate computation back into the residual stream format, essentially summarizing the complex high-dimensional processing into a form that can be added back to the information highway.

6. Transformer Blocks: The Core Architecture

The transformer block is the fundamental building unit, consisting of two key steps that always work together in a specific order:

1. Attention Mechanism: Routes and gathers relevant information from across the sequence (Communication Phase)
2. MLP Layer: Processes and transforms this contextualized information (Computation Phase)

This order is crucial: attention first allows each position to collect relevant context from the entire sequence, then the MLP can do informed processing with this enriched representation. The alternative (MLP → Attention) would force the MLP to process information blindly at each position before sharing, like trying to analyze the word "bank" without knowing whether the context involves rivers or money. By putting attention first, we follow the principle of "gather information, then process it" rather than "process locally, then share." This makes computation more effective - the MLP gets to do its complex reasoning on contextualized, relevant information rather than wasting computational effort processing isolated local information that may lack important context.

Each step is preceded by layer normalization (a standard preprocessing necessity before feeding information to any layer, including the final unembed step). Both the attention and MLP outputs are added back to the residual stream via residual connections.

This transformer block pattern repeats 12 times in GPT-2, with each iteration building increasingly complex representations from the combination of attention-driven information gathering and MLP-driven computation.

7. Unembed: Converting to Vocabulary Scores

The final step converts the high-dimensional representations into raw scores over the vocabulary. This is conceptually the reverse of embedding - we apply a linear transformation (matrix multiplication plus bias) from the model dimension (768) to vocabulary size (~50,000), creating one logit score for each possible token type in the vocabulary.

8. Logits: Raw Predictions

These raw scores (shown at the top of Fig. 2) represent the model's unprocessed predictions. During inference, they are converted to probabilities using softmax (not shown in Fig. 2 as it's a mathematical operation, not an architectural component), giving us probability estimates for each possible next token. For token generation, we can sample from this distribution or take the highest probability token. These probabilities also serve as the primary metric for measuring prediction quality, used for cross-entropy loss during both training and evaluation, perplexity calculations, and other language modeling metrics.

Potential Research Directions

Group Theory for Neural Circuits

Could we apply group theory (successfully used to organize the "particle zoo" in the 1950s-60s) to systematically group neural circuit patterns for better interpretability? Just as group theory revealed underlying symmetries that organized seemingly random hadrons into the Eightfold Way and eventually led to the quark model, similar mathematical frameworks might help classify the multitude of neural network pathways and patterns.

Many-Body Neural Systems

While we understand individual neurons, understanding group behavior remains elusive - reminiscent of the reductionist vs. holistic challenge in physics. Statistical mechanics deals with many-body systems and macroscopic properties. Could these approaches better explain neural network behavior than trying to understand individual components?

Scaling Laws Meet Phase Transitions

The sharp, discontinuous emergence of new abilities at critical model scales resembles phase transitions in physics, where smooth parameter changes trigger sudden qualitative shifts. This apparent paradox with smooth scaling laws mirrors how underlying variables (like temperature) change gradually while observable properties (like magnetization) change abruptly at critical points. Could identifying the neural equivalents of critical mass phenomena help predict emergence thresholds, or do we need frameworks that unify both gradual scaling improvements and sharp qualitative transitions in AI systems?

Methodological Questions

Measurement-Driven Discovery

I've learned that many breakthroughs in brain science emerged from defining precise measurable quantities. Karl Friston's free energy principle exemplifies this pattern—by mathematically defining free energy as an information-theoretic quantity, he unified previously separate theories under one measurable framework [4]. Similarly, Geoffrey Hinton's work at University of Toronto applied Boltzmann distributions from statistical mechanics to create breakthrough AI architectures, recently recognized with the 2024 Nobel Prize in Physics [5]. This pattern suggests we should think more systematically about what transformative measurable quantities might be waiting to be defined in AI interpretability research.

Measurement Strategy

Also, from an experimental physicist's perspective, what we measure is critical. While cross-entropy loss is optimized for training, is it actually the best way to understand model behavior? Should we measure raw token probabilities, perplexity, or calibration metrics instead of loss? Beyond probability measures, should we also focus on attention patterns, activation magnitudes, or gradient flows? How do we systematically measure emergent behavior with scaling? With so many measurable variables, how can we efficiently extract the high-level patterns that capture the most meaningful information?

Bayesian Approaches to Interpretability

Physics emphasizes precise measurements with quantified uncertainties, while industry ML often defaults to deterministic outputs. Given that the human brain operates as a Bayesian inference engine—continuously updating beliefs from limited evidence—should interpretability research embrace more explicitly Bayesian frameworks? This could be particularly valuable for understanding how transformers handle infrequent but critical decisions where we can't gather sufficient statistics, similar to how humans make consequential life choices without the luxury of controlled trials.

Fundamental Puzzles

Grounded vs. Abstract Meaning

Neural networks feel more mathematical than physical - in physics, information has specific meaning (e.g., particle spin states in wave functions correspond to measurable angular momentum), but in ML, any numerical representation is valid as long as it's useful for the task. Does certain neuron firing patterns have inherent meaning? This fundamental question about the nature of representation needs addressing.

In-Context Learning

How does transformer in-context learning compare to human rapid adaptation? The ability to change behavior without parameter updates seems more like pattern recognition than true learning—what can brain science tell us about this distinction?

Thoughts Under Fermentation

Residual Stream Design Choice

The linear additivity of residual streams makes interpretability possible. This seems to be a good choice, but I wonder if there are any other structural ways to allow interpretability more easily with different design setups.