# It Ain’t That Bad: Understanding the Mysterious Performance Drop in OOD
Generalization for Generative Transformer Models

Xingcheng Xu1, Zihao Pan2, Haipeng Zhang2, Yanqing Yang1,3 Corresponding
author.

###### Abstract

Generative Transformer-based models have achieved remarkable proficiency on
solving diverse problems. However, their generalization ability is not fully
understood and not always satisfying. Researchers take basic mathematical
tasks like $n$-digit addition or multiplication as important perspectives for
investigating their generalization behaviors. Curiously, it is observed that
when training on $n$-digit operations (e.g., additions) in which both input
operands are $n$-digit in length, models generalize successfully on unseen
$n$-digit inputs (in-distribution (ID) generalization), but fail miserably and
mysteriously on longer, unseen cases (out-of-distribution (OOD)
generalization). Studies try to bridge this gap with workarounds such as
modifying position embedding, fine-tuning, and priming with more extensive or
instructive data. However, without addressing the essential mechanism, there
is hardly any guarantee regarding the robustness of these solutions. We bring
this unexplained performance drop into attention and ask whether it is purely
from random errors. Here we turn to the mechanistic line of research which has
notable successes in model interpretability. We discover that the strong ID
generalization stems from structured representations, while behind the
unsatisfying OOD performance, the models still exhibit clear learned algebraic
structures. Specifically, these models map unseen OOD inputs to outputs with
equivalence relations in the ID domain. These highlight the potential of the
models to carry useful information for improved generalization.

## Introduction

Transformer-based models have exhibited remarkable advancements across diverse
domains, prominently in natural language processing (NLP) (Vaswani et al.
2017; Brown et al. 2020; Chowdhery et al. 2022) and fields such as computer
vision (Dosovitskiy et al. 2020) and biology (Jumper et al. 2021). Among
these, the extensive category of large language models, typified by instances
like ChatGPT (Ouyang et al. 2022) and GPT-4 (OpenAI 2023), has demonstrated
exceptional versatility, showcasing profound efficacy in tackling a myriad of
tasks, ranging from natural language challenges to code translation,
mathematical reasoning, and more (Bubeck et al. 2023; Trummer 2022; Zong and
Krishnamachari 2023). Although these accomplishments are undoubtedly
impressive, the generalization ability of these models is not fully understood
and not always satisfactory in issues such as semantic search (Asher et al.
2023), linguistic meaning capture (Bansal and Sharma 2023) and natural
language understanding (Bender et al. 2021).

Given the complexity of natural language tasks and the black-box nature of
these models, researchers view basic mathematical tasks such as $n$-digit
addition or multiplication as valuable avenues for gaining insights into their
generalization behaviors (Lee et al. 2023; Anil et al. 2022). Among them, many
have observed an interesting yet mysterious phenomenon when training on
$n$-digit operations (Brown et al. 2020; Anil et al. 2022; Jelassi et al.
2023). In cases where both input operands are $n$-digit long, the models
demonstrate excellent generalization on unseen $n$-digit inputs. However, they
unexpectedly and miserably struggle when faced with longer, unseen cases
(inputs with more than $n$ digits). For instance, when trained with operations
like $349+705=1054$, the model would perform well on unseen input $350+705$.
But when the inputs are $1349+2705$ which are longer in digits, the model
gives a wrong answer. This creates a clear distinction between the former,
known as in-distribution (ID) generalization, and the latter, termed out-of-
distribution (OOD) generalization.

Seeking to bridge this generalization gap, scholars have undertaken various
efforts to enhance OOD generalization. The techniques employed in this pursuit
encompass a diverse spectrum, including modifying position embeddings (Jelassi
et al. 2023) and attention mechanisms (Dubois et al. 2019), fine-tuning using
extended data samples, prompting and Scratchpad (Anil et al. 2022), priming
through selective longer-length data (Jelassi et al. 2023), and even utilizing
chain-of-thought (CoT) style data (Lee et al. 2023).

In spite of these different techniques, there is still a lack of understanding
regarding the underlying mechanism. The proposed solutions may therefore have
questionable robustness and become vulnerable to circumstance changes (Jelassi
et al. 2023). Considering the evident and notably poor OOD performance, it is
natural to ask whether it stems solely from random errors or if there is
anything informative learned by these models.

In this paper, we bring the mystery into attention and seek from the
mechanistic perspective (Nanda and Lieberum 2022; Zhong et al. 2023) in model
interpretability. This avenue of study offers a macroscopic understanding of
how neural networks work and has helped identify and interpret significant
phenomena such as “grokking”, also known as delayed generalization where
models exhibit improved generalization long after over-fitting their training
set (Liu et al. 2022).

Through training a set of lightweight generative Transformer models, including
NanoGPT and MinGPT (Karpathy 2022), on $n$-digit addition and multiplication
tasks, we have made an intriguing discovery. We find that the strong ID
generalization stems from structured representations, while the models have
learned a clear algebraic structure behind the unsatisfying OOD performance.
Specifically, these models map unseen OOD inputs to outputs with equivalence
relations in the ID domain. The representation learning process plays a
crucial role in facilitating both ID and OOD generalization observed in these
models. Initially, the representations are random. But as training progresses,
the structure of the learned representations becomes increasingly refined,
eventually allowing the models to accurately encode every input in the ID
domain. Concurrently, these structured representations are continuously
extended to map the unseen OOD domain. However, this extension does not occur
as ideally anticipated, resulting in the poor OOD performance. Thus, the
representation learning enables powerful ID generalization, but the continuous
extrapolation of these representations to OOD inputs gives rise to systematic,
rather than random, errors. The mechanistic insights from the discovered
patterns also highlight the potential of these models to make use of the
information for better generalization.

Our main contributions are as follows:

  * •

Showcasing the power of mechanistic empirical evaluation for LM
generalization: We train lightweight generative Transformer models (e.g.
NanoGPT, MinGPT) on arithmetic tasks to directly probe ID vs OOD
generalization111Our code is available at
https://github.com/xingchengxu/ExploreGPT, instead of resorting to
workarounds. Consequently, our approach offers macroscopic insights.

  * •

Discovering learned structure for OOD generalization: The discernible
algebraic structure would hopefully guide robust essential solutions for
strong OOD generalization.

  * •

Understanding the role of representations in generalization: We show that
representation learning enables strong ID performance, while unanticipated
extension of representations to OOD inputs leads to systematic errors.

## Related Work

### Generalization of Transformer-based Models in Arithmetic

Various studies have examined the performance of Transformer-based language
models in tasks involving arithmetic operations. Brown et al. (2020)
investigated the ability of GPT-3 to perform basic arithmetic operations
without task-specific training. Nogueira, Jiang, and Lin (2021) explored the
limitations of transformers in handling simple arithmetic operations.
Subsequent studies have further explored the generalization capabilities of
language models in arithmetic tasks. Qian et al. (2022) discovered that
language models exhibit poor OOD generalization, and traditional methods such
as explicit positional markers and fine-grained computation steps do not
effectively address this issue. To enhance the generalization ability of the
model, certain studies have approached the issue starting from a microscopic
perspective. For instance, Jelassi et al. (2023) replaced absolute position
embeddings with relative position embeddings. Additionally, Dubois et al.
(2019) suggested that utilizing a location-based attention mechanism proves
effective in the Lookup Table task. Other research has focused on the
intermediate learning process of the model. Anil et al. (2022) observed that
requesting the model to generate intermediate arithmetic steps before
providing the final output can improve generalization. Jelassi et al. (2023)
arrived at similar conclusions by decomposing the arithmetic pipeline and
improving generalization in five-digit addition tasks. In contrast, Lee et al.
(2023) presented a different perspective, emphasizing the importance of high-
quality, instructive data that can quickly elicit arithmetic capabilities.

While previous studies have primarily focused on evaluating or improving the
generalization capabilities of language models, our work has a different
objective, we aim to uncover the underlying mechanisms that govern
generalization. This explanatory goal, which seeks to understand the
foundations of generalization, has not been explicitly addressed in prior
research.

### Mechanistic Interpretability

Neural network interpretation has seen numerous studies focusing on various
types of models, including deep neural networks (DNNs) (Nam et al. 2020;
Barbiero et al. 2022), convolutional neural networks (CNNs) (Yuan et al. 2019;
Akhtar and Ragavendran 2020), and graph neural networks (GNNs) (Yuan et al.
2020; Xuanyuan et al. 2023). These works demonstrate diverse microscopic
interpretation techniques tailored to different architectures. From a
macroscopic perspective, Liu et al. (2022) tackle delayed generalization or
“grokking” using addition and modular addition tasks. They provide intuitive
explanations using effective theories and phase diagrams. Similarly, Zhong et
al. (2023) use modular addition to mechanistically explain algorithm discovery
in neural networks. Our work contributes to this growing field of mechanistic
interpretability by offering a macroscopic explanation specifically for
generative Transformer models.

## Preliminary and Experimental Setup

### Model Details

We employ the model framework of GPT, a Transformer with a decoder-only
architecture comprising multiple layers and multi-head attentions. We train
several small-scale models, namely NanoGPT and MinGPT (Karpathy 2022), from
random initialization using character-level tokenization and the conventional
next-token prediction objective. The training is conducted on basic
mathematical operations, specifically addition and multiplication of integers.
Detailed hyperparameters are shown in Table 1.

Hyperparameter |  Addition |  Multiplication
---|---|---
num layer |  3 |  6
num head |  3 |  6
dim embd |  48 |  192
vocab size |  10 |  10
context window |  15 |  19
dropout prob |  0.1 |  0.1
optimizer |  AdamW |  AdamW
learning rate |  0.0005 |  0.0005
betas |  (0.9, 0.95) |  (0.9, 0.95)
weight decay |  0.1 |  0.1
grad norm clip |  1.0 |  1.0
Table 1: Hyperparameter Information

### Dataset

The dataset is structured as a concatenation of operand pairs in a natural
order, with the reversed order of the operation results. This format,
demonstrated to be more conducive for learning in next-token prediction models
(Lee et al. 2023), offers a more approachable learning process. For instance,
consider the 3-digit addition $a+b=c$, represented as
“$a_{2}a_{1}a_{0}+b_{2}b_{1}b_{0}=c_{3}c_{2}c_{1}c_{0}$” in the standard
format. By reversing the output order of “$c$”, we obtain the reversed data
format “$a_{2}a_{1}a_{0}+b_{2}b_{1}b_{0}=c_{0}c_{1}c_{2}c_{3}$”. As we train
addition and multiplication models as distinct entities, we omit both the
operation symbols, i.e., $+$ and $\times$, and the equal sign, i.e., $=$, from
the dataset. Subsequently, the data undergoes character-level tokenization,
resulting in a vocabulary size of 10, corresponding to digits from 0 to 9.
When the context window surpasses the requisite size for a 3-digit addition,
we pad zeros before numbers “$a$”, “$b$”, and “$c$”. For instance, in the case
of 3-digit addition with a context window of 15, the addition expression
“$349+705=1054$” will be encoded as “$0034900705450100$”.

The dataset is partitioned into three distinct subsets: the training set
$\mathcal{D}_{1}$, randomly sampled from $n$-digit operations; the test set
$\mathcal{D}_{2}$, also drawn from $n$-digit operations but intentionally
devoid of any overlap with the training set (termed as the ID test set); and
an additional test set $\mathcal{D}_{3}$, sampled from $m$-digit operations
with $m>n$, where the value at positions greater than $n$ is non-zero
(referred to as the OOD test set).

In the experiments, we set $n=3$ and $m=5$ for both addition and
multiplication operations. Subsequently, from each of the datasets
$\mathcal{D}_{1}$, $\mathcal{D}_{2}$, and $\mathcal{D}_{3}$, we select 10,000
data points as the training set for addition and 50,000 for multiplication. We
sample 10,000 for the ID test set and OOD test set, respectively for both
operations.

### ID and OOD Domains

The data space is compartmentalized into three non-overlapping regions:
$\mathcal{D}_{1},\mathcal{D}_{2},$ and $\mathcal{D}_{3}$. The union of
$\mathcal{D}_{1}$ and $\mathcal{D}_{2}$ constitutes an ID domain, whereas
$\mathcal{D}_{3}$ represents an OOD domain. The models learn a function

$f:\mathcal{D}_{1}\cup\mathcal{D}_{2}\cup\mathcal{D}_{3}\rightarrow\mathcal{S},$

where $\mathcal{S}$ could be the output operation result space $\mathbb{N}$,
the output probability space, or the learned representation space.

Since the model is exclusively trained on the $\mathcal{D}_{1}$ space, the
acquired knowledge concerning $\mathcal{D}_{2}$ and $\mathcal{D}_{3}$ is an
extension of$\mathcal{D}_{1}$, albeit with an unclear underlying structure.
This constitutes the core aspect we seek to understand.

### Equivalence Classes

When training for addition and multiplication on $n$-digit operations, we have
identified a discernible algebraic structure. This is encapsulated in the
definition of the equivalence class $[(a,b)]_{p}$ for modular $p$, which is
elucidated as follows:

$[(a,b)]_{p}:=\\{(x,y)\in\mathbb{N}^{2}|\ x\equiv a\,(\textrm{mod}\,p),\
y\equiv b(\textrm{mod}\,p)\\}.$

The ensemble of these equivalence classes is denoted as

$\mathbf{Z}_{p}^{2}=\mathbf{Z}_{p}\times\mathbf{Z}_{p}=\\{[(a,b)]_{p}|\
(a,b)\in\mathbb{N}^{2}\\},$

where $\mathbf{Z}_{p}=\mathbb{Z}/p\mathbb{Z}$ on non-negative integers.

To illustrate, when training a model on 3-digit addition, we observe that the
learned operation function $f_{op}:\mathbb{N}^{2}\to\mathbb{N}$ effectively
translates to $f_{op}(a,b)=f_{op}([(a,b)]_{10^{3}})$, which will be stated in
the result section.

In alternative training data scenarios, the definition of equivalence classes
necessitates adaptation to accommodate specific contexts.

## Results

In this section, we present the key results and findings from our experiments.
These include observations on the phenomenon of generalization exhibited by
the models, the learned algebraic structure, as well as the probability and
representation structures in the model’s learning process.

### Generalization in OOD Domain

Figure 1: Training curves in addition and multiplication operations. Figure 2:
Contour plots for addition and multiplication operations.

Figure 1 depicts the training, ID test, and OOD test accuracy for addition and
multiplication operations in domains $\mathcal{D}_{1},\mathcal{D}_{2}$, and
$\mathcal{D}_{3}$ across different iterations. Panel (a) displays the training
curve for addition learned by NanoGPT, while Panel (b) showcases the curve for
multiplication learned by MinGPT. The hyperparameters employed by NanoGPT and
MinGPT can be found in Table 1.

Figure 3:  The probability distribution of each digit of the sequence in an
addition operation $c=a+b$. The left side of the black dashed line represents
the input $a+b$, while the right side is the result $c$. Figure 3(a) and
Figure 3(e) represent the $349+705$ and $128+256$, and the outputs are $1,054$
and $384$ ($450100$ and $483000$ in actual sequence output), respectively. In
the second column, we perturb the thousands digit of $a$: Figure 3(b)
represents $1,349+705$, and Figure 3(f) represents $3,128+256$. In the third
column, we perturb the thousands digit of $b$: Figure 3(c) represents
$349+2,705$, and Figure 3(g) represents $128+4,256$. In the fourth column, we
simultaneously perturb the thousands digit of $a$ and $b$: Figure 3(d)
represents $1,349+2,705$, and Figure 3(h) represents $3,128+4,256$.

By examining the figure, it becomes evident that both addition and
multiplication quickly converge to a stable state, achieving (almost) 100%
accuracy in training and ID testing in $\mathcal{D}_{1}$ and
$\mathcal{D}_{2}$. However, throughout the entire training process, the OOD
test accuracy remains zero for both 3-digit addition and multiplication in
$\mathcal{D}_{3}$. These results align with the discoveries made by Jelassi et
al. (2023) and Lee et al. (2023). When training on $n$-digit operations with
$n$-digit operands, the models demonstrate excellent generalization on unseen
$n$-digit inputs. Yet, they perform abysmally and mysteriously on longer,
unseen cases, establishing a contrast between ID generalization and OOD
generalization. Given the strikingly poor OOD performance, it is natural to
question whether it solely stems from random errors or if there is any
meaningful knowledge learned. The solution to the problem will be presented in
the subsequent subsection.

### Algebraic Structure

The mysterious absence of generalizability in the OOD domain prompts us to
delve deeper into the results. We begin by examining some samples from domains
$\mathcal{D}_{2}$ and $\mathcal{D}_{3}$. These examples are illustrated in
Table 2. When observing the 3-digit addition and multiplication cases, we
notice that the trained models produce incorrect results for the 4-digit
instances. Strikingly, these erroneous outputs mirror the results obtained
from the 3-digit cases. It appears that the model’s outputs peculiarly
“disregard” the thousands digit of the input numbers, irrespective of whether
it is an addition or multiplication operation.

Operands |  Output Result |  Correct Result
---|---|---
349 + 705 |  1,054 |  1,054
1,349 + 2,705 |  1,054 |  4,054
128 $\times$ 256 |  32,768 |  32,768
3,128 $\times$ 4,256 |  32,768 |  13,312,768
Table 2: Examples on models’ outputs for addition and multiplication.

To systematically analyze the behavior in the OOD domain $\mathcal{D}_{3}$, we
explore the entire two-dimensional lattice of 4-digit integers, namely
$\mathbb{N}^{2}\cap[0,10^{4})^{2}$. Figure 2 presents the contour plots for
the ground truth results of addition operation $c=a+b$ (Panel (a)) and
multiplication operation $c=a\cdot b$ (Panel (d)), with the number $a$ on the
horizontal axis and the number $b$ on the vertical axis. These landscapes
represent the expected learning and generalization capabilities of the models
on this lattice space.

However, when we utilize our trained models to generate results based on
3-digit operations, an unmistakably distinct pattern emerges, as depicted in
Panel (c) for addition and Panel (f) for multiplication. This prompts us to
investigate what structure the models have learned. We discover that there is
a modular relationship between the operands $a$ and $b$. The learned structure
can be represented as
$c=(a\,\,\textrm{mod}\,\,10^{3})\circ(b\,\,\textrm{mod}\,\,10^{3})$, where
$\circ$ represents either addition $+$ or multiplication $\times$. The ground
truth landscapes of these functions on the 4-digit integer lattice are
exhibited in Panel (b) for addition and Panel (e) for multiplication.
Visually, these two panels are identical to Panel (c) and Panel (f),
respectively. We compare the results of the operation
$(a\,\,\textrm{mod}\,\,10^{3})\circ(b\,\,\textrm{mod}\,\,10^{3})$ with the
outputs produced by the model. Surprisingly, they are identical across the
entire space $\mathbb{N}^{2}\cap[0,10^{4})^{2}$.

To formalize the results, we recall the definition of equivalence classes
$[(a,b)]_{p}$ for modular $p=10^{3}$:

$[(a,b)]_{p}:=\\{(x,y)\in\mathbb{N}^{2}|\ x\equiv a\,(\textrm{mod}\,p),\
y\equiv b(\textrm{mod}\,p)\\}.$

As $[(a,b)]_{p}$ is an equivalence class, we use the element in
$\mathbb{N}^{2}\cap[0,10^{3})^{2}=\mathcal{D}_{1}\cup\mathcal{D}_{2}$ to serve
as the representative of the class. The ensemble of these equivalence classes
then forms the space

$\mathbf{Z}_{p}^{2}=\mathbf{Z}_{p}\times\mathbf{Z}_{p}=\\{[(a,b)]_{p}|\
(a,b)\in\mathbb{N}^{2}\\},$

where
$\mathbf{Z}_{p}=\mathbb{Z}/p\mathbb{Z}=\\{[0]_{p},[1]_{p},\cdots,[p-1]_{p}\\}$.

The models trained on 3-digit addition and multiplication actually learned the
operation functions $f_{op}:\mathbf{Z}_{p}\times\mathbf{Z}_{p}\to\mathbb{N}$
for all integer paris on $\mathbb{N}\times\mathbb{N}$ such that
$f_{op}(a,b)=f_{op}([(a,b)]_{p})$ with $p=10^{3}$. As an example,
$f_{op}(1349,2705)=f_{op}([(349,705)]_{10^{3}})$. For addition, the learned
operation is $f_{+}(1349,2705)=f_{+}([(349,705)]_{10^{3}})=1054$, while the
learned multiplication is
$f_{\times}(1349,2705)=f_{\times}([(349,705)]_{10^{3}})=246045$.

As a summary of the results, the models learn to generalize the input in the
OOD domain $\mathcal{D}_{3}$ by assimilating equivalence classes in the ID
domain $\mathcal{D}_{1}\cup\mathcal{D}_{2}$. This result shows the limitations
of the models. However, this capability allows the models to extend their
learned knowledge beyond the ID domain $\mathcal{D}_{1}\cup\mathcal{D}_{2}$
shaped by the specific training data $\mathcal{D}_{1}$. Even though the output
is wrong, it is not so bad. They have still managed to acquire useful
information and demonstrate some level of learning.

In order to gain a deeper understanding of the training process for
Transformer models, we examine their token-level mapping using addition as an
example. Consider two ($n$+1)-digit numbers, where $a=a_{n}\times
10^{n}+\cdots a_{1}\times 10+a_{0}$ and $b=b_{n}\times 10^{n}+\cdots
b_{1}\times 10+b_{0}$. When training a Transformer model using randomly
sampled ($n$+1)-digit numbers, the model learns an approximate mapping from
the token-level input to the true function $c=a+b=c_{n+1}\times
10^{n+1}+\cdots c_{1}\times 10+c_{0}$. The learned approximation allows the
model to perform classification for each digit of the resulting sum $c$, as
follows:

$f_{Trans}(a_{n},\cdots,a_{0},b_{n},\cdots,b_{0})\approx(c_{0},c_{1},\cdots,c_{n+1}).$

However, if the highest digit is completely absent from the training data and
is instead padded with zeros, the training only guarantees learning of low
$n$-digit addition. In other words:

$f_{Trans}(0,a_{n-1},\cdots,a_{0},0,b_{n-1},\cdots,b_{0})\approx(c_{0},c_{1},\cdots,c_{n},0).$

This limitation may explain why it is challenging to generalize to higher
digits when the model is trained solely on lower digits. The absence of
examples with higher digits restricts the model’s ability to accurately
predict and generalize beyond the low-digit addition it has been trained on.

Building upon the observed algebraic structure discussed in the previous
context of this subsection, we also know that when testing the Transformer
models on higher digits that are non-zero, they do not significantly impact
the classification of each digit of $c$.

As a remark, it is important to note that when dealing with alternative
training data scenarios, the definition of equivalence classes may need to be
adjusted accordingly. For example, if the training data consists exclusively
of 1 and 3-digit operations, while OOD testing involves 2 and 4-digit numbers,
or if the training data includes 3-digit numbers for operand $a$ and 4-digit
numbers for operand $b$, the equivalence classes would require redefining to
account for these specific contexts.

### Probability Structure

Figure 4: 3D representation structure of the first three principal components
in the addition operation. Figure 4(a) to Figure 4(d) represent the initial
model, model with 14%, 51%, and 100% test accuracy, respectively.

In the preceding subsection, we examined the structured algebraic patterns
present in the output results. Considering that generative Transformer models
generate outputs based on probability distributions, our model employs a
greedy approach to select the output sequence with the maximum probability. We
now shift our focus from algebraic structures to the probability distribution
of the output sequences. Our objective is to investigate the underlying
factors that contribute to the emergence of these structured algebraic
patterns.

Specifically, we take two examples of 3-digit addition, namely $a+b$. We
introduce perturbations to the thousands digit of both $a$ and $b$, enabling
us to compare the variations in probability distributions before and after the
perturbations occur. This comparative study will shed light on the mechanisms
underlying the observed structured algebraic patterns.

Figure 3 displays the probability distributions of the next tokens in the
vocabulary at each position within the sequence for two addition examples:
$349+705$ and $128+256$. The plot showcases the probabilities before and after
perturbations for each token. Remarkably, we observe that regardless of
whether we perturb $a$ and $b$ separately or simultaneously, the probability
distribution in the model’s output sequence remains largely unchanged.

Furthermore, we note that the digits with the highest probability in the
output sequence remain consistent. This result implies that the algebraic
structure of the model expands from $\mathbf{Z}_{p}\times\mathbf{Z}_{p}$ to
$\mathbb{N}\times\mathbb{N}$. This expansion elucidates the structured
patterns depicted in Figure 2. Additionally, we conducted a systematic
examination of the entire integer lattice within
$\mathbb{N}^{2}\cap[0,10^{4})^{2}$. Notably, the results obtained from this
comprehensive analysis exhibit robustness, further supporting our findings.

### Representation Structure

Within the probability structure, we made a significant observation that the
model’s output remains insensitive to perturbations in the thousands digit.
This probability structure is rooted in the representation of the input
sequence, which can be expressed as follows:

$\mathbf{P}=\mathrm{Softmax}(\mathbf{WX}),$

where $\mathbf{P}\in[0,1]^{V\times L_{input}}$ represents the probability
matrix for the next tokens at each position, $\mathbf{W}\in\mathbb{R}^{V\times
d_{model}}$ signifies the learned weight matrix, and
$\mathbf{X}\in\mathbb{R}^{d_{model}\times L_{input}}$ denotes the learned
representation matrix of the input. The variables $V$, $L_{input}$, and
$d_{model}$ correspond to the vocabulary size, input length, and model
embedding dimension, respectively.

In this subsection, we delve deeper into the influence of these
representations on the probability structure, thereby shedding light on their
role in shaping the observed algebraic properties.

In order to explore the representations of $a+b$ in a systematic manner, we
conducted a thorough analysis on the two-dimensional integer lattice of
4-digit numbers. Specifically, for each input sequence $a+b$, we obtained a
high-dimensional embedding by considering the last column of the learned
representation matrix $\mathbf{X}$. Subsequently, we applied principle
component analysis (PCA) to project these embeddings into three or four
dimensions.

Figure 4 showcases the four different phases of representation observed during
the learning process of the model. The visualizations in the figure depict the
representations using the first three principle components. More specifically,
Figure 4(a) to 4(d) correspond to the random initial model, the model with
approximately 14%, 51%, and 100% ID test accuracy, respectively. The colors in
each figure correspond to the true units digit of the resulting $a+b$.

The observations made from Figure 4 demonstrate that the representations
gradually transition from disorderly to structured throughout the learning
process. Initially, the representations appear random with colors mixed
together (Figure 4(a)). However, as the training progresses, the structure of
the learned representations becomes increasingly refined (Figure 4(b) and
(c)), ultimately leading to the development of a well-learned representation
(Figure 4(d)) where each color is separated according to its true label.

To gain insight into the specific structures learned by the model, we focused
on the final phase of representation. We selected the first four principal
components of the representation using PCA, which accounted for an impressive
98.4% explained variance ratio. The contour plots of these principal
components are displayed in Figure 5. Figure 5(a) to (d) showcase the first to
fourth principle components of representations on the 2-dimensional integer
lattice of 4-digit numbers, respectively.

Figure 5: Representation structure of the first four principal components in
the addition operation. Figure 5(a) to Figure 5(d) represent the first to the
fourth principal components of model with 100% test accuracy, respectively.

Upon examining Figure 5, it becomes apparent that clear patterns emerge on
each graph. Furthermore, we observe that the representation structure on the
entire 2-dimensional integer lattice of 4-digit numbers is similar to that
observed on the lattice of 3-digit numbers. These patterns align with the
algebraic structures that have been previously identified. Importantly, even
when perturbing the thousands digit of the input, the representation of the
input remains robust. This robustness demonstrates the model’s resilience to
perturbations and reinforces the stability of the probability distribution.

### From Representation to Algebraic Structure

The findings discussed above regarding algebraic structures, probability
distributions, and representation structures also hold true for multiplication
operations.

The systematic analysis approach outlined earlier provides a comprehensive
understanding of the model’s generalization capabilities through the
assimilation of equivalence classes present within the ID domain. The
representation structures successfully incorporate the assimilation of these
equivalence classes, thereby extending the ID structure to OOD scenarios via
the probability distribution of sequences. Consequently, this assimilation
becomes evident within the algebraic structures as well.

## Discussion

We study the generalization behaviors of lightweight generative Transformer
models, specifically in the context of arithmetic tasks. In this section, we
discuss some limitations of our investigation.

Firstly, our study focused exclusively on arithmetic tasks, such as $n$-digit
addition and multiplication problems. While arithmetic serves as a useful
domain for analyzing generalization due to its structured nature, the findings
may not directly translate to other types of NLP tasks or domains.

Secondly, defining ID and OOD domains for natural language is challenging.
Unlike arithmetic, where the numerical values and operations define the ID
domain, NLP involves a wide range of linguistic structures, vocabularies, and
contextual dependencies. The boundaries between ID and OOD examples become
blurred, making it difficult to precisely delineate the domains.

Moreover, determining equivalence relations in NLP tasks can be considerably
more complicated. As demonstrated in the paper, equivalence relations allow
the models to map inputs based on shared characteristics or properties and
thus plays a key role in the generalization behaviors observed in arithmetic.
However, due to the inherent ambiguity and subjectivity of language, much more
efforts may be needed to establish clear-cut equivalence relations in NLP
tasks.

## Conclusion

In conclusion, our study investigates the generalization behaviors of
generative Transformer models trained on arithmetic tasks. Through empirical
evaluation, we make several key discoveries that provide insight into their
generalization abilities.

Our findings reveal that these models demonstrate strong generalization within
the trained distribution, performing well on arithmetic problems within this
domain. This can be attributed to the structured representations learned
during training, which progressively refine to accurately encode ID inputs.
However, our investigation also uncovers an underlying algebraic structure
that contributes to the models’ unsatisfactory performance on OOD inputs. The
models attempt to map OOD inputs using equivalence relations within the ID
domain, leading to errors and a lack of robustness in OOD scenarios. The
representation learning process plays a crucial role in enabling both ID and
OOD generalization. While representations become refined for ID inputs,
extending them to OOD inputs does not occur optimally, resulting in suboptimal
performance.

Despite challenges in OOD generalization, our findings suggest that these
models hold valuable information for improved generalization. The structured
representations offer crucial insights that, if utilized effectively, could
enhance the models’ capabilities and improve OOD performance. To advance
generalization capabilities further, novel strategies could be explored to
address the issues raised in this study, including optimizing representation
extension, reducing systematic errors, and promoting robust performance.

## References

  * Akhtar and Ragavendran (2020) Akhtar, N.; and Ragavendran, U. 2020.  Interpretation of intelligence in CNN-pooling processes: a methodological survey.  _Neural computing and applications_ , 32(3): 879–898.
  * Anil et al. (2022) Anil, C.; Wu, Y.; Andreassen, A.; Lewkowycz, A.; Misra, V.; Ramasesh, V.; Slone, A.; Gur-Ari, G.; Dyer, E.; and Neyshabur, B. 2022.  Exploring length generalization in large language models.  _Advances in Neural Information Processing Systems_ , 35: 38546–38556.
  * Asher et al. (2023) Asher, N.; Bhar, S.; Chaturvedi, A.; Hunter, J.; and Paul, S. 2023.  Limits for Learning with Language Models.  _arXiv preprint arXiv:2306.12213_.
  * Bansal and Sharma (2023) Bansal, P.; and Sharma, A. 2023.  Large Language Models as Annotators: Enhancing Generalization of NLP Models at Minimal Cost.  _arXiv preprint arXiv:2306.15766_.
  * Barbiero et al. (2022) Barbiero, P.; Ciravegna, G.; Giannini, F.; Lió, P.; Gori, M.; and Melacci, S. 2022.  Entropy-based logic explanations of neural networks.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 36, 6046–6054.
  * Bender et al. (2021) Bender, E. M.; Gebru, T.; McMillan-Major, A.; and Shmitchell, S. 2021.  On the dangers of stochastic parrots: Can language models be too big?  In _Proceedings of the 2021 ACM conference on fairness, accountability, and transparency_ , 610–623.
  * Brown et al. (2020) Brown, T.; Mann, B.; Ryder, N.; Subbiah, M.; Kaplan, J. D.; Dhariwal, P.; Neelakantan, A.; Shyam, P.; Sastry, G.; Askell, A.; et al. 2020.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33: 1877–1901.
  * Bubeck et al. (2023) Bubeck, S.; Chandrasekaran, V.; Eldan, R.; Gehrke, J.; Horvitz, E.; Kamar, E.; Lee, P.; Lee, Y. T.; Li, Y.; Lundberg, S.; et al. 2023.  Sparks of artificial general intelligence: Early experiments with GPT-4.  _arXiv preprint arXiv:2303.12712_.
  * Chowdhery et al. (2022) Chowdhery, A.; Narang, S.; Devlin, J.; Bosma, M.; Mishra, G.; Roberts, A.; Barham, P.; Chung, H. W.; Sutton, C.; Gehrmann, S.; et al. 2022.  PALM: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_.
  * Dosovitskiy et al. (2020) Dosovitskiy, A.; Beyer, L.; Kolesnikov, A.; Weissenborn, D.; Zhai, X.; Unterthiner, T.; Dehghani, M.; Minderer, M.; Heigold, G.; Gelly, S.; et al. 2020\.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_.
  * Dubois et al. (2019) Dubois, Y.; Dagan, G.; Hupkes, D.; and Bruni, E. 2019.  Location attention for extrapolation to longer sequences.  _arXiv preprint arXiv:1911.03872_.
  * Jelassi et al. (2023) Jelassi, S.; d’Ascoli, S.; Domingo-Enrich, C.; Wu, Y.; Li, Y.; and Charton, F. 2023\.  Length Generalization in Arithmetic Transformers.  _arXiv preprint arXiv:2306.15400_.
  * Jumper et al. (2021) Jumper, J.; Evans, R.; Pritzel, A.; Green, T.; Figurnov, M.; Ronneberger, O.; Tunyasuvunakool, K.; Bates, R.; Žídek, A.; Potapenko, A.; et al. 2021\.  Highly accurate protein structure prediction with AlphaFold.  _Nature_ , 596(7873): 583–589.
  * Karpathy (2022) Karpathy, A. 2022.  A minimal PyTorch re-implementation of the OpenAI GPT (Generative Pretrained Transformer) training.  _GitHub https://github.com/karpathy/minGPT_.
  * Lee et al. (2023) Lee, N.; Sreenivasan, K.; Lee, J. D.; Lee, K.; and Papailiopoulos, D. 2023.  Teaching Arithmetic to Small Transformers.  _arXiv preprint arXiv:2307.03381_.
  * Liu et al. (2022) Liu, Z.; Kitouni, O.; Nolte, N. S.; Michaud, E.; Tegmark, M.; and Williams, M. 2022\.  Towards understanding grokking: An effective theory of representation learning.  _Advances in Neural Information Processing Systems_ , 35: 34651–34663.
  * Nam et al. (2020) Nam, W.-J.; Gur, S.; Choi, J.; Wolf, L.; and Lee, S.-W. 2020.  Relative attributing propagation: Interpreting the comparative contributions of individual units in deep neural networks.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 34, 2501–2508.
  * Nanda and Lieberum (2022) Nanda, N.; and Lieberum, T. 2022.  A mechanistic interpretability analysis of grokking.  In _Alignment Forum_.
  * Nogueira, Jiang, and Lin (2021) Nogueira, R.; Jiang, Z.; and Lin, J. 2021.  Investigating the limitations of transformers with simple arithmetic tasks.  _arXiv preprint arXiv:2102.13019_.
  * OpenAI (2023) OpenAI. 2023.  GPT-4 Technical Report.  _ArXiv_ , abs/2303.08774.
  * Ouyang et al. (2022) Ouyang, L.; Wu, J.; Jiang, X.; Almeida, D.; Wainwright, C.; Mishkin, P.; Zhang, C.; Agarwal, S.; Slama, K.; Ray, A.; et al. 2022.  Training language models to follow instructions with human feedback.  _Advances in Neural Information Processing Systems_ , 35: 27730–27744.
  * Qian et al. (2022) Qian, J.; Wang, H.; Li, Z.; Li, S.; and Yan, X. 2022.  Limitations of language models in arithmetic and symbolic induction.  _arXiv preprint arXiv:2208.05051_.
  * Trummer (2022) Trummer, I. 2022.  CodexDB: Synthesizing code for query processing from natural language instructions using GPT-3 Codex.  _Proceedings of the VLDB Endowment_ , 15(11): 2921–2928.
  * Vaswani et al. (2017) Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A. N.; Kaiser, Ł.; and Polosukhin, I. 2017.  Attention is all you need.  _Advances in neural information processing systems_ , 30.
  * Xuanyuan et al. (2023) Xuanyuan, H.; Barbiero, P.; Georgiev, D.; Magister, L. C.; and Lió, P. 2023\.  Global concept-based interpretability for graph neural networks via neuron analysis.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 37, 10675–10683.
  * Yuan et al. (2019) Yuan, H.; Chen, Y.; Hu, X.; and Ji, S. 2019.  Interpreting deep models for text analysis via optimization and regularization methods.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, 5717–5724.
  * Yuan et al. (2020) Yuan, H.; Tang, J.; Hu, X.; and Ji, S. 2020.  Xgnn: Towards model-level explanations of graph neural networks.  In _Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining_, 430–438.
  * Zhong et al. (2023) Zhong, Z.; Liu, Z.; Tegmark, M.; and Andreas, J. 2023.  The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networks.  _arXiv preprint arXiv:2306.17844_.
  * Zong and Krishnamachari (2023) Zong, M.; and Krishnamachari, B. 2023.  Solving math word problems concerning systems of equations with GPT-3.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 37, 15972–15979.