PaperlayerView exact revision on arXiv

Source-linked AI summary

Analyzing Transformers in Embedding Space

Guy Dar, Mor Geva, Ankit Gupta, Jonathan Berant

arXiv:2209.02535v3cs.CLcs.LG

TL;DR

Transformer interpretability usually requires input passes, while prior zero-pass methods covered only some parameters or simplified architectures. This paper projects all Transformer parameters into embedding space, supporting parameter interpretation, cross-model alignment, and zero-shot classifier stitching, with 70% IMDB accuracy in 30% of cases without training.

  • Problem

    Existing Transformer interpretation commonly requires input-dependent forward or backward passes, while earlier zero-pass approaches covered only selected parameters or simplified architectures.

  • Method

    The paper projects Transformer parameters into embedding space and uses a right-inverse to express model inner products as vocabulary-item interactions.

  • Results

    30% of stitching cases reached 70% accuracy on IMDB without training after transferring a fine-tuned classifier through embedding spaces.

  • Takeaways & Limitations

    Embedding space can support parameter interpretation and comparison or transfer across Transformer models that share a vocabulary.

  • Takeaways & Limitations

    The framework requires a linear language-modeling head, focuses mainly on GPT-2, and does not rule out interpretive factors beyond the vocabulary lens.

Abstract

from arXiv · show

Understanding Transformer-based models has attracted significant attention, as they lie at the heart of recent technological advances across machine learning. While most interpretability methods rely on running models over inputs, recent work has shown that a zero-pass approach, where parameters are interpreted directly without a forward/backward pass is feasible for some Transformer parameters, and for two-layer attention networks. In this work, we present a theoretical analysis where all parameters of a trained Transformer are interpreted by projecting them into the embedding space, that is, the space of vocabulary items they operate on. We derive a simple theoretical framework to support our arguments and provide ample evidence for its validity. First, an empirical analysis showing that parameters of both pretrained and fine-tuned models can be interpreted in embedding space. Second, we present two applications of our framework: (a) aligning the parameters of different models that share a vocabulary, and (b) constructing a classifier without training by ``translating'' the parameters of a fine-tuned classifier to parameters of a different model that was only pretrained. Overall, our findings open the door to interpretation methods that, at least in part, abstract away from model specifics and operate in the embedding space only.

1 Introduction

The paper introduces a zero-pass framework that interprets all Transformer parameters in embedding space and uses that shared space to compare or transfer information across models. Evidence covers pretrained and fine-tuned models, including alignment across BERT instances and untrained classifier stitching.

  • 1 Introduction: The approach is motivated by the cost and input dependence of activation-based interpretation, while requiring a linear language-modeling head and focusing mainly on GPT-2.The authors state that the framework requires a linear language-modeling head.
  • 1 Introduction: The framework interprets all pretrained language-model weights in embedding space, extending earlier work beyond feed-forward values and simplified attention networks.It covers feed-forward keys and values together with all attention parameters.
  • 1 Introduction: The method projects model quantities into embedding space and back with a right-inverse, recasting model inner products as vocabulary-item interactions.This applies to query-key interactions and value-output interactions.
  • 1 Introduction: Different BERT instances pretrained from distinct random seeds aligned layer i with layer i, indicating convergence to semantically similar solutions.The alignment compares models that share a vocabulary.
  • 1 Introduction: The framework provides qualitative and quantitative evidence that pretrained and fine-tuned Transformer parameters can be interpreted in embedding space.Fine-tuning GPT-2 for movie-review sentiment produced vocabulary items characterizing sentiment toward movies.
  • 1 Introduction: 30% of stitching cases reached 70% accuracy on IMDB without training after transferring a fine-tuned sentiment classifier through embedding spaces.The transferred classifier came from a fine-tuned model, while the target model was pretrained only.

2 Background

The Transformer background describes layers composed of attention and feed-forward modules operating on sequences of vectors, with residual connections carrying information across layers. Attention forms masked interactions among projected queries, keys, and values before combining head outputs.

  • 2 Background: Each Transformer layer contains an attention module followed by a position-wise two-layer feed-forward module, with residual connections combining their outputs.The input and output sequences contain N vectors of dimensionality d.
  • 2 Background: Attention projects sequence representations into queries, keys, and values, splits them across H heads, and computes one masked attention map per head.The mask M has shape N×N, and each map is applied to its corresponding value head.
  • 2 Background: The head outputs are concatenated and projected through the output matrix before the residual connection adds the module input.The supplied architecture description identifies concatenation, output projection, and residual addition.

X + Concat

The paper views Transformer computation through residual streams and represents hidden states and parameter interactions in embedding space. Attention interaction matrices describe token-pair attention and how attended tokens affect subsequent hidden states.

  • X + Concat: The feed-forward module uses learnable key and value matrices and can be written as f(QK^T)V, with its output added through a residual connection.Keys are weights of the first layer and values are weights of the second layer.
  • X + Concat: The embedding matrix maps vocabulary items into model representations and can project final-layer hidden states back into vocabulary space.The same matrix is often reused for input embeddings and output projection in architectures including GPT-2.
  • X + Concat: Residual connections carry a hidden-state stream through layers, while modules read from it and write information back into it.The final hidden state can be projected into embedding space to interpret layerwise predictions.
  • X + Concat: Attention parameters are reorganized into WQK and WVO interaction matrices that summarize query-key attention and value-to-output transitions.WQK encodes attention between token pairs, while WVO describes how attending to tokens affects later hidden states.
  • X + Concat: The attention computation can therefore be restated using interaction matrices, separating attention maps from the value transformations applied by each head.The attention map is expressed in terms of WQK, while head outputs use the corresponding value-output interaction.

3 Parameter Projection

The paper projects Transformer parameters into embedding space, recasting model operations as vocabulary-level interactions. This framework covers feed-forward and attention components, using an approximate transpose-based projection in practice.

  • Transformer parameters can be projected into embedding space to support interpretation of their behavior.
  • The framework uses a right-inverse E′ of the embedding matrix to recast model inner products as interactions between vocabulary items.
  • Attention Module: Attention outputs are represented as embedding-space transitions, while attention subheads expose interpretable relations between input and output vocabulary items.
  • FF Module: Feed-forward keys act as concept detectors that control the coefficients of corresponding value vectors, linking their embedding-space interpretations.
  • In practice, the method uses E′ = E^T rather than an exact pseudo-inverse because interpretation retains only the top-k vocabulary items.

4 Interpretability Experiments

Experiments examine whether projected parameters correspond to dynamically computed hidden states and whether embedding-space projections reveal semantic structure. They find interpretable relations in attention and feed-forward parameters, hidden-state alignment, and fine-tuning changes.

  • 4.1 Parameter Interpretation Examples: Attention transition matrices reveal relations involving word parts, gender, geography, orthography, part-of-speech tags, and semantic topics.
  • 4.1 Parameter Interpretation Examples: Feed-forward key-value pairs with at least 15% top-100 overlap comprise approximately 5% of pairs and often concern related topics.
  • 4.1 Parameter Interpretation Examples: Embedding-space lookup retrieves feed-forward keys or values related to programming, measurements, and animals from averaged topic embeddings.
  • 4.2 Hidden State and Parameters: Activated parameter vectors mirror hidden-state representations to some extent, as measured by higher aligned-state than randomly sampled-state Rk comparisons.
  • 4.3 Interpretation of Fine-tuned Models: Fine-tuning changes in WQ, WK, WV, and K project to sentiment-related words, whereas changes in V and WO do not show this behavior.

5 Aligning Models in Embedding Space

The embedding space supports comparing parameters across models that share a vocabulary, enabling layer alignment and zero-shot stitching. BERT layers trained from different seeds align by index, while stitching transfers a fine-tuned classifier to a pretrained model with inconsistent but sometimes above-random IMDB accuracy.

  • 5 Aligning Models in Embedding Space: The method compares parameter vectors across models by projecting them into a shared embedding space when the models use the same vocabulary and tokenizer.
  • 5.1 Layer Alignment: Parameters from corresponding BERT layers tend to align across models and parameter groups, whereas feature-space matching appears random and unintelligible.
  • 5.2 Zero-shot Stitching: A stitching kernel projects hidden states from model A through embedding space into model B, where they feed B’s fine-tuned layers without additional training.
  • 5.2 Zero-shot Stitching: Stitching can exceed random accuracy on IMDB but does not do so consistently across random seeds.
  • 5.2 Zero-shot Stitching: Three of 11 stitched models reached roughly 70% IMDB accuracy, significantly above the 50% random baseline, when stitching over top layers.

6 Related Work

Related work has analyzed hidden representations and individual Transformer components, often using input-dependent activations. This paper instead studies static parameters comprehensively and unifies prior embedding-space analyses across Transformer components.

  • 7 Related Work: Prior interpretability research mainly probes hidden representations and tracks their evolution across layers.
  • 7 Related Work: Other studies analyze individual neurons, attention heads, or feed-forward values, mostly through input-dependent activations.
  • 7 Related Work: This work extends two-layer attention analysis and feed-forward-value interpretation into a unified embedding-space framework for all Transformer components.

7 Discussion

The framework offers a low-compute, input-independent perspective on Transformer parameters and places parameters from different models in a canonical embedding space. The paper demonstrates model alignment and stitching as consequences of this representation.

  • 7 Discussion: Input-independent interpretation may reveal intrinsic parameter properties while reducing dependence on particular inputs.
  • 7 Discussion: The method requires only a fraction of typical compute and memory and permits interpreting individual parameters without input-dependent passes.
  • 7 Discussion: A canonical embedding space enables model-agnostic parameter comparison across models and supports alignment and stitching applications.

8 Limitations

The framework is limited by its vocabulary-centered perspective, the chosen embedding-matrix variant, and concentration of experiments on GPT-2. The authors also note possible stereotype-related associations in parameters.

  • 8 Limitations: The vocabulary lens does not rule out other factors, and the chosen E′ = E^T projection may be improved by future variants.
  • 8 Limitations: Most experiments focus on GPT-2 because of current framework shortcomings and presentation considerations, although BERT suggests nonlinearities may be resolvable.
  • 8 Limitations: Parameters may associate related terms because of stereotypes learned from the training corpus.

A Rethinking Interpretation

The section evaluates embedding-space inverses for interpreting Transformer parameters and shows that the transpose embedding matrix is effective for top-k interpretation at small k, with limitations at larger k and in GPT-2 large.

  • Inverse choice: 0.83 versus 0.10: E^T better matches inner products than E+ under top-k reconstruction for 300 random vectors at k = 1000.The scores compare E^T and the pseudo-inverse E+ using cosine similarity between original and reconstructed inner products.
  • Inverse choice: E^T outperforms E+ for small k across normal, hidden-state, and FF-value distributions, while E+ approaches E^T only for hidden states.The comparison applies to GPT-2 medium, the paper’s main model, and uses k values from 10 through 500.
  • Inverse choice: At larger k, E^T becomes increasingly worse as an approximate right-inverse, except that hidden states retain high performance.The higher-k analysis reverses the small-k trend for the tested distributions.
  • Model variation: GPT-2 base gives similar conclusions, whereas GPT-2 large shows zigzagging E+ behavior and is usually better with E+, though E^T is more predictable.The authors leave GPT-2 large’s counter-intuitive zigzag behavior for future work.
  • Evaluation: The metric compares overlap between the top-k activated token sets after projecting vectors into embedding space.The top-k sets correspond to the most activated vocabulary items, and their overlap is measured with a Jaccard index.

B.2 Final Prediction and Parameters

Embedding-space parameter interpretations correlate with model predictions and reveal relationships among corresponding parameter groups across layers and models.

  • Final prediction and parameters: Final predictions correlate with the most activated parameter vectors from each layer, with stronger correspondence for aligned than randomized hidden states.The comparison covers FF keys, FF values, attention values, and attention outputs.
  • Corresponding parameter pairs: Matching parameter pairs show higher embedding-space similarity than shuffled pairs for FF, attention value-output, and query-key groups.Figure 6 evaluates aligned versus shuffled pairs across layers using average Simk at k = 100.
  • Examples: The supplied examples include GPT-2 medium, GPT-2 base, and GPT-2 large attention heads, with corpus-vocabulary filtering marked by double asterisks.The examples manually select a subset of GPT-2 medium’s 384 attention heads.

C.1.1 Low-Level Language Modeling

The section presents qualitative examples of low-level language-modeling patterns recovered from attention-head transition pairs across GPT-2 models.

  • GPT-2 medium examples: GPT-2 attention heads yield token transitions reflecting recurring relations such as sources-to-ources, tales-to-stories, and journals-to-magazines.These examples come from GPT-2 medium Layer 21 Head 7 and illustrate related lexical or subword transitions.
  • GPT-2 medium examples: GPT-2 medium Layer 19 Head 13 captures word-boundary relations, pairing initial fragments or consonants with word-final tokens in examples such as R-to-senal for arsenal.The listed pairs include both directions of several fragment-to-word relationships.
  • Cross-model examples: GPT-2 large examples include where-to-upon, with-to-regard, and feminine-pronoun-to-name associations.The listed heads are GPT-2 large Layer 27 Heads 6 and 12.
  • Cross-model examples: Additional heads show geographic and nationality associations, including Saudi-to-Arabia, Finland-to-onen, and India-to-Mumbai.These examples are drawn from GPT-2 base and GPT-2 medium heads.

C.3 Feedforward Keys and Values

The section illustrates feedforward key-value pairs whose top-k vocabulary overlaps identify interpretable topics and concepts across layers and dimensions.

  • Selection criterion: Key-value pairs are retained when at least 15% of their top-k vocabulary items overlap at k = 100.The section calls the value index within a layer a dimension.
  • Selection criterion: The examples exclude tokens outside the corpus vocabulary when marked with double asterisks.This filtering convention is explicitly noted for the displayed lists.

D Sentiment Analysis Fine-Tuning Vector Examples

The section visualizes classifier fine-tuning vectors by sentiment label and samples parameter-group differences across layers 9–11. The displayed vectors contain recurring lexical patterns that distinguish positive- and negative-associated terms.

  • Fine-tuning vectors are labeled by sentiment, with separate vectors corresponding to the POSITIVE and NEGATIVE classifier labels.The section introduces the classifier-weight difference visualization using these two sentiment directions.
  • The POSITIVE and NEGATIVE vectors contain contrasting lexical associations, including favorable terms on one side and unfavorable terms on the other.The displayed examples include words associated with sentiment polarity, though the extracted lists are noisy and fragmented.
  • The analysis samples four difference vectors for each parameter group across fine-tuned layers 9–11, retaining examples with manually observed relevant patterns.Parameter groups include feed-forward keys and values plus attention query, key, value, and output subheads.
  • Additional difference-vector examples show recurring clusters involving sentiment, quality, media, policy, and other vocabulary categories across sampled parameter groups.These examples include both favorable and unfavorable lexical patterns in different sampled vectors.
  • The remaining sampled vectors likewise mix positive, negative, and topical vocabulary, illustrating heterogeneous patterns across parameter groups rather than a single uniform direction.Examples include sentiment terms alongside media, regulatory, financial, and general-language associations.
  • Later vectors continue this pattern with mixed sentiment and topical terms, including media, finance, regulation, and general vocabulary.The displayed samples extend the qualitative inspection across further parameter-group differences.
Loading 2209.02535v3…