Table of Contents
Fetching ...

Weight-based Analysis of Detokenization in Language Models: Understanding the First Stage of Inference Without Inference

Go Kamoda, Benjamin Heinzerling, Tatsuro Inaba, Keito Kudo, Keisuke Sakaguchi, Kentaro Inui

TL;DR

The paper presents a weight-based decomposition of GPT-2's first attention layer to explain detokenization without performing any forward passes. By separating token- and position-derived components, it defines six terms $T^{\text{ee}}$, $T^{\text{pp}}$, $T^{\text{pe}}$, $T^{\text{ep}}$, $T^{\text{e}}$, and $T^{\text{p}}$, and demonstrates that $T^{\text{ee}}$ captures token affinity while $T^{\text{p}}$ and $T^{\text{pp}}$ bias attention toward nearby tokens. The work links detokenization to absolute position embeddings, LayerNorm variance, and token frequency via $T^{\text{e}}$, and provides empirical validation on OpenWebText with insights into head-specific roles. While limited to the first layer and GPT-2, the findings offer a principled, weight-based mechanism for early inference in LMs and point to broader applicability across architectures and position-embedding schemes.

Abstract

According to the stages-of-inference hypothesis, early layers of language models map their subword-tokenized input, which does not necessarily correspond to a linguistically meaningful segmentation, to more meaningful representations that form the model's "inner vocabulary". Prior analysis of this detokenization stage has predominantly relied on probing and interventions such as path patching, which involve selecting particular inputs, choosing a subset of components that will be patched, and then observing changes in model behavior. Here, we show that several important aspects of the detokenization stage can be understood purely by analyzing model weights, without performing any model inference steps. Specifically, we introduce an analytical decomposition of first-layer attention in GPT-2. Our decomposition yields interpretable terms that quantify the relative contributions of position-related, token-related, and mixed effects. By focusing on terms in this decomposition, we discover weight-based explanations of attention bias toward close tokens and attention for detokenization.

Weight-based Analysis of Detokenization in Language Models: Understanding the First Stage of Inference Without Inference

TL;DR

The paper presents a weight-based decomposition of GPT-2's first attention layer to explain detokenization without performing any forward passes. By separating token- and position-derived components, it defines six terms , , , , , and , and demonstrates that captures token affinity while and bias attention toward nearby tokens. The work links detokenization to absolute position embeddings, LayerNorm variance, and token frequency via , and provides empirical validation on OpenWebText with insights into head-specific roles. While limited to the first layer and GPT-2, the findings offer a principled, weight-based mechanism for early inference in LMs and point to broader applicability across architectures and position-embedding schemes.

Abstract

According to the stages-of-inference hypothesis, early layers of language models map their subword-tokenized input, which does not necessarily correspond to a linguistically meaningful segmentation, to more meaningful representations that form the model's "inner vocabulary". Prior analysis of this detokenization stage has predominantly relied on probing and interventions such as path patching, which involve selecting particular inputs, choosing a subset of components that will be patched, and then observing changes in model behavior. Here, we show that several important aspects of the detokenization stage can be understood purely by analyzing model weights, without performing any model inference steps. Specifically, we introduce an analytical decomposition of first-layer attention in GPT-2. Our decomposition yields interpretable terms that quantify the relative contributions of position-related, token-related, and mixed effects. By focusing on terms in this decomposition, we discover weight-based explanations of attention bias toward close tokens and attention for detokenization.
Paper Structure (34 sections, 24 equations, 13 figures, 1 table)

This paper contains 34 sections, 24 equations, 13 figures, 1 table.

Figures (13)

  • Figure 1: Focusing on the token/position embeddings, first LayerNorm layer, and the first attention layer, we conduct weight analyses and show high attention scores are assigned to close (top middle; \ref{['sec:term-pos-self', 'sec:term-pos-compare', 'sec:term-positions']}) and related (top right; \ref{['sec:detokenization-token']}), supporting the detokenization hypothesis. We also show that the high attention score to the first token derives from LayerNorm (bottom middle; \ref{['sec:wpe-var', 'sec:high-pos-var-at-first']}). Regarding token embedding, we also discuss the relationship between token frequency and LayerNorm (bottom right; \ref{['sec:wte-variance']}).
  • Figure 2: A: Examples of support for detokenization. When the current position token is "iens", the past token that yields the largest $T^{\text{ee}}_{}$ value (=Rank 1) is "_sap" in head#4 and head#7. B: Heatmap of $T^{\text{ee}}_{}$ for head#7 and head#1. Tokens are randomly sampled from the vocabulary for visualization. C: ROC of head#7 and head#1 when token $i$ is "iens". D: Average AUROC for each head. Heads with high AUROC values contribute to the reconstruction of bi-grams, consequently contributing to detokenization.
  • Figure 3: A: $T^{\text{p}}_{j}$ for all context token position $j$ for head#1 (top) and #7 (bottom). The shaded area represents the variance of the term deriving from $e_{\text{ID}_i}$ Y-axis titles in this, and subsequent figures are inset for readability. B: $T^{\text{pp}}_{500, j}$$={\bm{p}_i\bm{W}_h^{QK}\bm{p}_j^\top}/{\sigma_i\sigma_j}$ for context token position $j\leq 500$. The blue, green, and orange lines show scores with mean, maximum, and minimum standard deviation for $\sigma_j$: $\sigma_j = \frac{1}{|V|}\sum_{\text{ID}}\sigma(\bm{e}_{\text{ID}} + \bm{p}_j)$, $\sigma_j = \text{max}_{\text{ID}} \sigma(\bm{e}_{\text{ID}} + \bm{p}_j)$, and $\sigma_j = \text{min}_{\text{ID}} \sigma(\bm{e}_{\text{ID}} + \bm{p}_j)$, respectively. C,D: Sum of $T^{\text{pp}}_{500, j}$ and $T^{\text{p}}_{j}$ and its result after taking softmax with a temperature of $\sqrt{d'}$. E: Empirical attention weights aggregated over texts in OpenWebText Corpus when present token position $i=500$ focusing on the last few $j$ positions. The red line and area show the empirically observed weights and the blue line corresponds with the blue lines in D.
  • Figure 4: Contribution of the 6 terms in \ref{['eq:attn_s_6terms']} for each current token position $i$ for head#1 and #7
  • Figure 5: Relation between token frequency and $T^{\text{e}}_{}$ for head#1 and head#7.
  • ...and 8 more figures