Source-linked AI summary
The Bottom-up Evolution of Representations in the Transformer: A Study with Machine Translation and Language Modeling Objectives
Elena Voita, Rico Sennrich, Ivan Titov
TL;DR
The paper asks how Transformer token representations evolve across layers under different learning objectives, addressing limitations of probing-based descriptions. It uses mutual information, canonical correlation analysis, and similarity comparisons to analyze information flow, finding distinct objective-specific patterns: LM loses past information, MLM separates context encoding from token prediction, and MT retains more word-type information.
Problem
Prior probing-task observations showed objective-dependent Transformer representations, but the processes producing these differences and their information flow remained insufficiently explained.
Method
The paper analyzes layerwise individual-token representations with mutual information estimators, canonical correlation analysis, and similarity-based property comparisons.
Results
LM loses past information while forming future predictions; MLM first encodes context, partially forgets token identity, then recreates it; MT refines representations while retaining most word-type information.
Takeaways & Limitations
The findings support viewing representation evolution as a compression/prediction trade-off and provide a possible explanation for MLM's superior pretraining performance over LM.
Takeaways & Limitations
Direct mutual-information estimation for finite samples in multidimensional spaces is challenging, so subsequent analyses use canonical correlation analysis and related frameworks.
Abstract
from arXiv · showhide
We seek to understand how the representations of individual tokens and the structure of the learned feature space evolve between layers in deep neural networks under different learning objectives. We focus on the Transformers for our analysis as they have been shown effective on various tasks, including machine translation (MT), standard left-to-right language models (LM) and masked language modeling (MLM). Previous work used black-box probing tasks to show that the representations learned by the Transformer differ significantly depending on the objective. In this work, we use canonical correlation analysis and mutual information estimators to study how information flows across Transformer layers and how this process depends on the choice of learning objective. For example, as you go from bottom to top layers, information about the past in left-to-right language models gets vanished and predictions about the future get formed. In contrast, for MLM, representations initially acquire information about the context around the token, partially forgetting the token identity and producing a more generalized token representation. The token identity then gets recreated at the top MLM layers.
1 Introduction
This paper characterizes how Transformer token representations evolve across layers under MT, LM, and MLM objectives, rather than evaluating them only through probing tasks. It finds objective-specific information flows, including compression and prediction stages that differ across tasks.
- Motivation and approach: The study explains layerwise representation behavior by characterizing information flow rather than relying only on auxiliary probing-task performance.The analysis focuses on how individual token representations evolve between layers under different learning objectives.
- Methods: The authors analyze token representations using mutual information, canonical correlation analysis, and similarity-based comparisons of token properties.They examine information about token identity, output labels, positions, neighboring tokens, and CCG supertags.
- Findings by objective: LM representations lose information about the past while forming predictions about the future from bottom to top layers.The authors argue that neither current-token information and its past nor future information is represented well in LMs.
- Findings by objective: MLM representations first encode surrounding context and partially forget token identity, then recreate token identity at the top layers.These stages correspond to context encoding followed by token prediction in the compression/prediction view.
- Findings by objective: MT representations are refined using context with less processing, retaining most word-type information throughout the layers.The paper presents this contrast as part of its explanation for why MLM may outperform LM for pretraining.
- Contributions: The paper proposes viewing layerwise token-representation evolution through a compression/prediction trade-off.Its experiments distinguish information loss about inputs from accumulation of information about outputs, especially for MLM.
2 Tasks
The paper compares Transformer models trained for machine translation, left-to-right language modeling, and masked language modeling. Each task defines different inputs and output labels, with MT predictions produced through the decoder and LM or MLM predictions from final token representations.
- Task definitions: The study defines input X and output Y for each of the three Transformer tasks.The tasks are machine translation, standard left-to-right language modeling, and masked language modeling.
- Machine translation: NMT predicts each target word from the source sentence and previously generated target words.The model estimates p(y_i|X, y_1,i−1, θ) for target positions.
- Task differences: For MT, top encoder representations encode information for the decoder rather than directly predicting output labels.This differs from LM and MLM, where output labels are predicted from final token representations.
- Language modeling: In language modeling, the model predicts x_t from preceding tokens, using the final representation of token x_t−1.The objective estimates P(x_t|x_1, . . . , x_t−1, θ).
- Masked language modeling: Masked language modeling randomly masks 15% of tokens and predicts the original token from the final representation of its corrupted version.The selected input is replaced by [MASK], a random token, or left unchanged according to the stated masking procedure.
3 Data and Setting
The experiments use comparable Transformer setups across objectives, with shared data, architecture, initialization, multiple datasets, and repeated random initializations. The data come from WMT English–German and English–French translation corpora, while language models use source-side text.
- Experimental controls: The study controls comparisons by using the same training data, architecture, and parameter initialization across tasks.Experiments are repeated with multiple datasets and different random initializations for reliability.
- Data: The experiments use WMT 2017 English–German with 5.8m sentence pairs and WMT 2014 English–French with 40.8m sentence pairs.The English–French language-modeling analysis uses source-side data after removing 2.8m sentence pairs, leaving 38m for training.
- Model setup: All models follow the Transformer base-model setup and training procedure.Implementation details are provided in the appendix.
4 The Information-Bottleneck Viewpoint
The Information Bottleneck viewpoint treats Transformer layers as an information-flow process whose patterns depend on the learning objective. Mutual-information analyses show objective-specific changes in token representations across layers for MT, LM, and MLM.
- Information Bottleneck principle: The Information Bottleneck principle views representations as retaining output-relevant information while compressing input information.The output variable determines which input features are relevant and should be preserved.
- Information flow in sequence models: This study models MT, LM, and MLM as functions from an input token sequence X to an objective-specific output Y, analyzing individual-token representations at every layer.The analysis focuses on how information about input or output tokens is retained as representations evolve.
- Mutual-information estimation: Mutual information is estimated by clustering token representations from frequent words across layers into N = 10000 clusters using mini-batch k-means with k = 100.The experiments gather 5 million occurrences of the 1000 most frequent (sub)words for each model and layer.
- Results: For input-token information, LM representations lose information monotonically, MT loses it less sharply, and MLM first loses then recovers it at its two top layers.The MLM phases are termed context encoding and token reconstruction, with the recovered information suggested to come from other context tokens.
- Results: When input and output tokens are compared, input information decreases while output information increases for both LM and MLM.For MLM, the two processes are more distinct and correspond to context encoding followed by token prediction; MT shows only encoding because nothing is predicted directly.
- Methodological caveat: Direct mutual-information estimation from finite samples in multidimensional spaces is challenging, motivating canonical correlation analysis for subsequent corroboration.The later analysis continues to ask how information flows across layers and what information each layer represents.
5 Analyzing Changes and Influences
The analysis uses PWCCA to compare representations across tasks, layers, and token properties, revealing objective-specific information flow and changes. LM and MLM show non-monotonic layer transformations, while MT changes decrease toward the top.
- 5.1 Canonical Correlation Analysis: PWCCA compares token representations across models or layers by maximizing linear correlation; similarity ranges from 0 to 1.A value of 1 indicates linear transformations of each other, while 0 indicates no correlation.
- 5.2.1 Distances between tasks: Training-objective differences are much larger than differences caused by random initialization.This suggests PWCCA captures differences in learned information rather than training randomness.
- 5.2.1 Distances between tasks: MT and MLM representations are closer to each other than either is to LM representations.MT and MLM use the whole sentence context and focus on reconstructing or translating the current token, unlike left-to-right LM.
- 5.2.1 Distances between tasks: MT representation change decreases monotonically from bottom to top layers, unlike the non-monotonic changes in LM and MLM.The authors relate the latter patterns to encoding and reconstruction phases.
- 5.3 Fine-grained analysis: Frequent tokens change more than rare tokens in LM and MT, but this frequency dependence fades near the top of LM.The authors speculate that top LM layers prioritize future prediction over incorporating the past.
- 5.3 Fine-grained analysis: MLM’s transition from generalized representations to reconstructed token identities requires more change for rare tokens.This makes the two MLM processing stages especially pronounced in the fine-grained analysis.
- 5.3.1 Varying token frequency: Rare tokens generally influence other tokens more than frequent tokens, especially at the first MT layer and across LM layers.MLM does not show this pattern at lower layers, and token dropout removes the extreme rare-token influence in MT and LM.
6 What does a layer represent?
The study evaluates which token properties are preserved in each layer by comparing nearest representations. Results show objective-specific retention of identity, position, context, and syntax.
- 6 What does a layer represent?: The analysis tests whether nearest token representations share identity, position, neighboring-token information, or CCG supertags.It uses token occurrences and checks whether each property matches among their closest representations.
- 6.2.1 Preserving token identity: MLM representations initially lose token identity and later recreate it at higher layers, while MT and LM occurrences of related words become mixed.For MLM, masked states form a separate first-layer cluster before becoming disambiguated across layers.
- 6.2.1 Preserving token identity: t-SNE tracks occurrences of related tokens across layers, with colors identifying token types and the horizontal axis representing layers.For the “is,” “are,” “was,” and “were” visualization, the cited results describe mixing in MT and LM and disambiguation in MLM.
- 6.2.1 Preserving token identity: Position information is dismissed more strongly in LM and MLM than in MT; MT retains noticeable positional ordering even at the last layer.The analysis interprets position as input information that may be less relevant to the output.
- 6.2.2 Preserving context: LM reduces the importance of previous-token information while forming information about future tokens; MLM neighbor importance rises through layer two and then declines.The MLM pattern is consistent with context encoding followed by token reconstruction.
- 6.2.3 Preserving syntactic information: CCG-tag information degrades at higher MLM layers, whereas LM first accumulates left-context syntax and later dismisses it while forming future-context information.The LM analysis separates CCG information corresponding to previous and next parts of the sentence.
7 Additional related work
The paper complements probing-based analyses by studying how representations evolve between layers and by applying CCA-based methods to task-dependent information flow.
- 7 Additional related work: Earlier work probed Transformer representations for morphology, syntax, and semantics, whereas this study analyzes how representations evolve between layers.The paper positions its approach as complementary to probing tasks.
- 7 Additional related work: The paper uses CCG supertags derived with the Yoshikawa et al. tagger in its syntactic representation analysis.The cited implementation passage identifies the tagger and its ELMO-based latest version.
- 7 Additional related work: CCA-based methods have been used to study learning dynamics, layer dimensionality, memorization, generalization, and important neurons in neural networks.Prior applications include CNNs, RNNs, and neural machine translation models.
8 Conclusions
The paper analyzes objective-dependent information flow through token representations using a compression/prediction perspective. Its experiments support this view and offer a possible explanation for MLM’s superior pretraining performance over LM.
- 8 Conclusions: The paper frames token-representation evolution as a compression/prediction trade-off determined by the learning objective.It supports this perspective through experiments across different objectives and relates the findings to probing-task observations.
A Data and Setting
The study uses Transformer-base configurations for machine translation and encoder-only language models, with shared tokenization and training procedures.
- Machine translation models use six encoder and six decoder layers, eight attention heads, dmodel = 512, and dff = 2048.
- Language models use only the Transformer encoder with the same hyperparameters as the translation models.
- Sentences use byte-pair encoding with approximately 32000-token source and target vocabularies.
- Training uses minibatches of approximately 15000 source tokens, approximate length batching, and the optimizer and learning-rate schedule of Vaswani et al. (2017).
B Fine-grained analysis of change and influence: varying PoS
The analysis compares token change and influence across parts of speech, finding patterns similar to those observed across frequency groups.
- Parts of speech containing frequent tokens generally change more and influence less.
- Figures 13 and 14 report token change and token influence, respectively, for different parts of speech.
C.1 Preserving token identity: experimental setup
The experiment tests whether models confuse token representations by comparing selected main tokens with nearby contrastive tokens in embedding space.
- For each selected main token, the experiment chooses contrastive tokens that could potentially be confused with it.
- The analysis uses part-of-speech tags from Stanford CoreNLP.
- Figures 13 and 14 provide related analyses of token change and influence by part of speech.
- Contrastive tokens are the top 10 closest tokens to each main token in each model’s embedding space.
- The study gathers 9000 occurrences of contrastive tokens and adds them to 1000 states of the corresponding main token.