Transformer Feed-Forward Layers Build Predictions by Promoting Concepts in the Vocabulary Space

In plain words

A language model guesses the next word by stacking many layers. Each layer adds a change to a running summary of the text so far. One kind of layer, the feed-forward block, makes these changes too.

The paper asks what those changes actually do. Its answer: a feed-forward block's change is a sum of many pieces. Each piece is a fixed direction, scaled by a number. If you read that direction as a list of words, it often lines up with an idea, such as breakfast foods or pronouns. Most of the time, the words that end up on top are the ones that these large pieces push up, rather than the ones that other pieces push down.

The authors use this in two ways. They turn up a few "safe" pieces to make a model less toxic. They also watch the big pieces to decide when a model can stop early and skip its remaining layers.

TL;DR

• The output of each feed-forward layer is a weighted sum of value vectors. Projecting a value vector onto the vocabulary often gives a readable concept §3 §4.
• Dominant sub-updates mostly promote the final top token, rather than eliminate rivals §5.1.
• Turning on 10 hand-picked non-toxic value vectors in GPT-2 reduces the share of completions with Perspective toxicity above 0.5 from 58.5% to 30.8% on a challenging prompt subset (author-reported) §6.1 tab. 5.
• A rule based on the top-10 dominant sub-updates exits early with 94.1% accuracy and saves 20.0% of layers on WikiLM (author-reported) §6.2 tab. 6.

Problem

Earlier work on how language models build predictions mostly followed hidden representations across layers, or probed the model with target tasks. The feed-forward (FFN) layers, and how their output feeds the next-token distribution, were less studied §1 §7.

The paper assumes each token representation can be read as a distribution over the vocabulary at any layer, and that each FFN output is an additive change to that distribution §2. It also assumes that the FFN output can be split into sub-updates, one per value vector, and that each sub-update can be read on its own §3. The concept analysis uses a 16-layer word-level WikiLM and a 12-layer GPT-2. The promotion analysis uses a 24-layer GPT-2 §4 §5.

Method

Setup §2. A transformer block adds the FFN output to the residual stream. The model's next-token distribution is a softmax of the output embedding applied to the final representation eq. 1. Reading a representation at an intermediate layer uses the same projection, so each layer gives a distribution over the vocabulary §2.

Decomposition §3. The FFN is a two-matrix block with a non-linearity between them. Its output is a sum over value vectors, each weighted by a coefficient that depends on the input. Each term, a coefficient times a value vector, is a sub-update §3. The effect of one sub-update on the probability of a token w is a product of two factors, so its effect on the logit is a dot product of the token embedding with the sub-update eq. 2.

Two readings of a sub-update §3. The first factor is static: the projection of the value vector onto the vocabulary ranks tokens independently of the input. The second factor, the coefficient, is the same for all tokens for a given input, so it compares how much each value vector contributes in that update §3.

Concept annotation §4. Experts looked at the top-30 tokens of each sampled value vector, looked for patterns that recur in at least 4 tokens, described each pattern, and labelled it semantic, syntactic or (WikiLM only) names §4 app. A.2. The authors compare these results with random vectors. Top tokens of real value vectors contain more concepts than random vectors do §4.1 tab. 2.

Promotion versus elimination §5.1. For each candidate token, the authors look at the scores of the 10 dominant sub-updates in a layer. A dominant sub-update is one with a large share of the layer output, where share is the absolute coefficient times the value vector norm, relative to the sum over all value vectors app. A.3. They collect two kinds of events:

1. Saturation: a layer whose update leaves the final token as the top candidate through the last layer.
2. Elimination: a layer whose update drops the current top candidate to a rank above 1.

For promoted tokens, the maximum score is large and the minimum is small negative. For eliminated tokens, the scores are small in both directions. Random sub-updates produce scores close to zero §5.1 tab. 4. The authors conclude that top tokens are promoted by dominant sub-updates §5.1.

Extreme sub-updates §5.2. Value vectors are clustered with agglomerative clustering on cosine distance, with k = 10^4 clusters for each model §5.2. Clusters that often produce extreme scores are of two kinds: upper-layer vectors that promote rare tokens, which the authors call saturation vectors, and vectors spread over all layers that promote common tokens such as stopwords. Together these cover about 1.1% (WikiLM) and 1.7% (GPT-2) of value vectors §5.2.

Toxicity suppression §6.1. The authors pick 10 value vectors whose top tokens are positive words such as "safe", "thank" and "transparency", either by manual search or by grading token lists with the Perspective API §6.1 app. A.4. At generation time these 10 coefficients are set to 3, and the model generates 20 tokens for each prompt §6.1.

Early exit §6.2. For each example and layer, the 10 dominant sub-updates are mapped to their clusters. A set of clusters seen at saturation events of a layer is compared with clusters seen without saturation at later layers, by intersection size. The model exits when the overlap with the saturation sets exceeds the overlap with the later non-saturation sets, both on average and at the maximum §6.2 app. A.5. The paper footnotes that this is a simplification of the split by saturation layer §6.2.

The three procedures, in pseudocode. The sub-update share and the promotion scores follow the paper's definitions §3 eq. 2 app. A.3. The steering and early-exit steps follow the paper and the authors' code.

# Split one FFN output into sub-updates  [§3]
# Input:  coefs (d_m,), values (d_m, d)
# Output: sub_updates (d_m, d), dominant (indices)

def split_ffn_output(coefs, values, top=10):
    # A sub-update is a coefficient times its value vector  [§3]
    sub_updates = coefs[:, None] * values              # (d_m, d)
    # Share of the output: |m| * ||v|| over the layer total  [app. A.3]
    size = abs(coefs) * norm(values, axis=1)           # (d_m,)
    share = size / sum(size)
    # The dominant sub-updates are the largest by share  [§5] [app. A.3]
    dominant = argsort(-share)[:top]
    return sub_updates, dominant

# Score dominant sub-updates for a candidate token w  [§5.1] [eq. 2]
# Input:  w (token id), coefs and values of the dominant sub-updates
# Output: one signed score per sub-update

def candidate_scores(w, coefs, values, embedding):
    # Score = embedding of w dotted with the sub-update  [eq. 2]
    # Positive promotes w, negative eliminates it  [§5.1]
    return [
        dot(embedding[w], c * v)
        for c, v in zip(coefs, values)
    ]

# Steer generation with fixed value vectors  [§6.1]
# Input:  prompt, chosen (dict: layer -> value indices), coef = 3
# Output: 20 generated tokens

def steer_generation(prompt, chosen, coef=3):
    # Set each chosen coefficient to coef at every position  [§6.1]
    # Code overwrites the output of the first FFN matrix, pre-GeLU
    # [code: toxic_language_suppression/toxic_suppression_wrapper.py#L147-L161]
    install_hooks(chosen, coef)
    # Generate 20 tokens; the script defaults to 3 beams, no sampling
    # [code: toxic_language_suppression/toxicity_scoring.py#L56-L62]
    return generate(prompt, new_tokens=20)

# Early exit from dominant clusters  [§6.2] [app. A.5]
# Input:  clusters (set of cluster ids of the top-10 sub-updates)
# pos[layer]: sets seen at a saturation event at that layer
# neg[layer]: sets seen without saturation at that layer

def should_exit(layer, clusters, pos, neg):
    # Overlap with this layer's saturation sets  [§6.2]
    own = [len(clusters & s) for s in pos[layer]]
    for later in range(layer + 1, num_layers):
        # Overlap with later non-saturation sets  [§6.2]
        other = [len(clusters & s) for s in neg[later]]
        # Exit only if mean and max overlap both beat the later sets
        # [code: early_exit/wikilm-earlyexit-ours.ipynb]
        if mean(own) <= mean(other) or max(own) <= max(other):
            return False
    return True

Hyperparameters

Cluster and early-exit notebook settings are cited from the code where the paper is silent.

Results

Author-reported. Setup: WikiLM and GPT-2 as above; RealToxicPrompts challenging subset with Perspective API scores, 20-token continuations; WikiLM for early exit §4 §6.1 §6.2.

Concept projections §4.1 tab. 2. Share of top-scoring tokens associated with a concept.

Promotion vs. elimination §5.1 tab. 4. Scores of the reference token from the 10 dominant sub-updates (GPT-2 with 24 layers; WikiLM 16 layers).

Toxicity suppression §6.1 tab. 5. GPT-2, challenging RealToxicPrompts subset, percentage of completions above 0.5 per Perspective attribute. Arrows give relative change from the GPT-2 baseline.

Early exit on WikiLM §6.2 tab. 6. Accuracy is the share of examples where exiting at the predicted layer gives the full model's prediction. Saved layers counts only correct predictions.

The paper reports these as "on-par" with the classifier baselines §6.2.

Paper vs code

Official code: github.com/aviclu/ffn-values, linked from the paper's abstract footnote. The repository's README lists the four authors and links Avi Caciularu's homepage, and the owner is that same account. It has no LICENSE file, so it has no reuse license; no code is quoted in this card. Commit read: 81ac29236c944dbcd7b1956d5728fcc4c48b05fc.

• LayerNorm in the projection. Main text: value vectors are projected without LayerNorm, and LayerNorm is omitted for GPT-2 §2 §2 footnote. Appendix: the findings also hold when LayerNorm is applied, with 64.5% top-30 overlap app. A.1. The GPT-2 wrapper projects with the final LayerNorm applied to the value row toxic_suppression_wrapper.py. The WikiLM projection is raw README.md. Which projection produced the 10 manual picks is not stated.
• Coefficient point of intervention. The paper defines the coefficient as the activation after the non-linearity, f(W_K x) §3. The steering hook overwrites the output of the first FFN matrix, before GeLU, for the chosen indices toxic_suppression_wrapper.py. Inferred from the code, not stated in the paper: the effective coefficient is GeLU(3), about 2.996, not exactly 3.
• Decoding. The paper states 20-token continuations and, in a commented-out line, beam search of size 3 §6.1. The script defaults to no sampling and 3 beams toxicity_scoring.py toxicity_scoring.py.
• Clustering. The §5.2 clustering (k = 10^4, cosine distance, agglomerative) is not in this repository. The README points to an external repository for it README.md. The early-exit notebook loads precomputed clusters and uses 200 of them, matching the appendix value k = 2e2 [code: early_exit/wikilm-earlyexit-ours.ipynb].
• Concept annotation and sub-update share. The expert annotations are not in the repository. I found no script for the share computation in the directories the README lists. The data-generation script was not read.
• Steering vectors. The 10 manual picks are stored in a pickle file in the repository non_toxic_values.pickle. Not verified against the appendix list, because the file was not parsed.

No other discrepancies found.

Pitfalls

• Model size for toxicity is unstated. §6.1 says "GPT2" without a size. The code defaults to gpt2-medium [code: toxic_language_suppression/toxic_suppression_wrapper.py#L10]. The appendix picks sit at layers 14 to 23 app. A.4, which a 12-layer model cannot have, so the toxicity model is not the 12-layer GPT-2 of §4. The paper does not name the checkpoint.
• Layer indexing. The paper writes value vectors as layer superscript and index subscript, e.g. layer 14, index 1853 app. A.4. Check whether the index is 0- or 1-based before you reuse the list. The code's indexing is not stated in the paper.
• Projection choice changes the concepts. The main text omits LayerNorm, but the GPT-2 wrapper includes it. Top-30 overlap is 64.5% on average, so about a third of the tokens differ app. A.1.
• Manual selection is not reproducible in code. The 10 picks came from a search of under five minutes, chosen greedily, with no optimisation app. A.4. The paper does not say whether the search used the RealToxicPrompts evaluation prompts. Treat the 47% figure as a single chosen set.
• The toxicity gain is partly paid for. Perplexity rises from 21.7 to 25.3, more than SD (23.9) tab. 5. The paper calls this relatively small.
• Abstract says "almost 50%". The toxicity attribute falls 47% tab. 5; other attributes fall 40 to 63%. Make sure the claim you cite names the attribute.
• Toxicity is reduced, not removed. The ethics statement says the method only lowers the probability of toxic language [Ethics Statement].
• Coefficient 3 was not swept. It was picked as "relatively high" by reading a figure §6.1. The paper reports no other values.
• Concept labels are manual and noisy. Annotating one value vector took 8.5 minutes on average, and the authors call the concept frequencies an underestimate §4. Inter-annotator agreement is not stated.
• Dominant sub-updates can be negative. With GeLU, coefficients may be negative. The share uses absolute values, but the promotion and elimination scores use signed values app. A.3 §5.1.
• Only the top-10 dominant sub-updates are analysed. They cover about 0.24% of value vectors and about 5% of the output app. A.3. Other sub-updates are ignored by the analysis, so the findings describe only these.
• Early-exit rule is a simplification. The paper footnotes that the rule splits the non-saturation sets by saturation layer §6.2. The notebook requires both mean and maximum intersections to be larger [code: early_exit/wikilm-earlyexit-ours.ipynb]. Implement the notebook, not the short description.
• Early-exit baselines use more data. The classifiers train on 8k examples with 8-fold tuning, while the rule uses 10k examples app. A.5. Compare them with the same split.
• Possible copy error in tab. 6. The rows for x and x̂ show identical numbers. The text does not explain this, so check the source before relying on either row.
• Clustering is an external dependency. Reproducing the cluster sets needs code outside this repository README.md.
• Word-level and sub-word models differ. WikiLM has 267,744 word tokens; its concept analysis and its early-exit results do not transfer to GPT-2 without a check §4 §6.2.
• Figures are not reproduced here. Figs. 2 and 3 are described in words only fig. 2 fig. 3.

Lineage

• Builds on: Geva et al. 2021, "Transformer Feed-Forward Layers Are Key-Value Memories", which reads FFN parameters in vocabulary space and treats them as key-value memories §3 §7. The paper extends that reading from whole layers to individual value vectors, and from upper layers to all layers §4.2.
• Compared against: Self-Debiasing (Schick et al., 2021) and WordFilter (RealToxicityPrompts, Gehman et al., 2020) for toxicity §6.1; layer-wise logistic-regression classifiers for early exit §6.2.
• Superseded by: not stated in this wiki.