Applied Science, Faculty of
Electrical and Computer Engineering, Department of
DSpace
UBCV
Mahdavi, Seyed Mohammad Sadegh
Master of Applied Science - MASc
University of British Columbia
In this thesis, we investigate the ability of neural networks, particularly Transformers, to reason and memorize. First, we focus on graph neural networks and Transformers, and analyze their performance on algorithmic reasoning tasks. We show that while models can achieve high accuracy on data from the same distribution as their training data, their performance drops significantly when faced with new, out-of-distribution data. We further show that even high performance on benchmark numbers may be misleading and true reasoning capability of these models remains limited. We identify several challenges involved in achieving true reasoning abilities and generalization to new data. We propose solutions to some of these challenges, including fixing input representation issues, hybrid models, and enlarging the training dataset. We also examine the expressivity of Transformers, providing a theoretical analysis of their ability to memorize data points. The results show a linear relationship between a Transformer's memory capacity and both the number of its attention heads as well as the input's context size.
Exploring Algorithmic Reasoning and Memorization inTransformers: Challenges and InsightsbySeyed Mohammad Sadegh MahdaviB.Sc. Computer Engineering, Sharif University of Technology, 2021A THESIS SUBMITTED IN PARTIAL FULFILLMENTOF THE REQUIREMENTS FOR THE DEGREE OFMaster of Applied ScienceinTHE FACULTY OF GRADUATE AND POSTDOCTORAL STUDIES(Electrical and Computer Engineering)The University of British Columbia(Vancouver)August 2023© Seyed Mohammad Sadegh Mahdavi, 2023The following individuals certify that they have read, and recommend to the Facultyof Graduate and Postdoctoral Studies for acceptance, the thesis entitled:Exploring Algorithmic Reasoning and Memorization in Transformers:Challenges and Insightssubmitted by Seyed Mohammad Sadegh Mahdavi in partial fulfillment of therequirements for the degree of Master of Applied Science in Electrical andComputer Engineering.Examining Committee:Renjie Liao, Assistant Professor, Electrical and Computer Engineering, UBCCo-SupervisorChristos Thrampoulidis, Assistant Professor, Electrical and Computer Engineering,UBCCo-SupervisorPurang Abolmaesumi, Professor, Electrical and Computer Engineering, UBCSupervisory Committee MemberNick Harvey, Professor, Computer Science, UBCSupervisory Committee MemberiiAbstractIn this thesis, we investigate the ability of neural networks, particularly Transformers,to reason and memorize. First, we focus on graph neural networks and Transformers,and analyze their performance on algorithmic reasoning tasks. We show that whilemodels can achieve high accuracy on data from the same distribution as theirtraining data, their performance drops significantly when faced with new, out-of-distribution data. We further show that even high performance on benchmarknumbers may be misleading and true reasoning capability of these models remainslimited. We identify several challenges involved in achieving true reasoning abilitiesand generalization to new data. We propose solutions to some of these challenges,including fixing input representation issues, hybrid models, and enlarging thetraining dataset. We also examine the expressivity of Transformers, providing atheoretical analysis of their ability to memorize data points. The results show alinear relationship between a Transformer’s memory capacity and both the numberof its attention heads as well as the input’s context size.iiiLay SummaryDeep Learning, particularly the Transformer architecture, has been successful invarious fields, including computer vision, healthcare, and language processing. Thisresearch investigates the ability of these models to reason and memorize. Whilemodels can perform well within their training data, their performance is poor andunpredictable when faced with new, out-of-distribution data. This suggests thatdespite their vast training, these models’ reasoning ability is limited. The researchalso examines the memorization capacity of Transformers, providing a theoreticalanalysis of their ability to remember data. The findings highlight the challenges inachieving true reasoning abilities in artificial intelligence models.ivPrefaceThis thesis is based on the following two works. The first work has been peer-reviewed and published. The second work is under review.• Sadegh Mahdavi, Kevin Swersky, Thomas Kipf, Milad Hashemi, ChristosThrampoulidis, and Renjie Liao. “Towards better out-of-distribution general-ization of neural algorithmic reasoning tasks.” In: Transactions on MachineLearning Research (TMLR), 2023• Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. “Memorizationcapacity of multi-head attention in transformers.” Preprint, 2023Design of the research work, conducting the research, and the analysis of resultsis done by the author.vTable of ContentsAbstract . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . iiiLay Summary . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ivPreface . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . vTable of Contents . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . viList of Tables . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ixList of Figures . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . xiAcknowledgments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . xiv1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11.1 Background . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 31.1.1 Graph Neural Networks . . . . . . . . . . . . . . . . . . 41.1.2 Transformers . . . . . . . . . . . . . . . . . . . . . . . . 41.1.3 Neural Algorithmic Reasoning . . . . . . . . . . . . . . . 51.1.4 Large Language Models . . . . . . . . . . . . . . . . . . 51.2 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . 61.2.1 Out-of-Distribution Generalization and Reasoning . . . . 61.2.2 In-Distribution Generalization . . . . . . . . . . . . . . . 71.2.3 Memorization . . . . . . . . . . . . . . . . . . . . . . . . 72 Out-of-Distribution Generalization for Neural Algorithmic Reasoning 9vi2.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 92.2 Background . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 112.3 Current Challenges of OOD generalization of Neural AlgorithmicReasoning Tasks . . . . . . . . . . . . . . . . . . . . . . . . . . 132.3.1 Difference with Common OOD Settings . . . . . . . . . . 142.3.2 Distribution Shift in Input Representations . . . . . . . . 162.3.3 Unrepresentative Dataset Generation . . . . . . . . . . . . 172.3.4 Model Selection is Hard with Nearly Perfect ID Performances 182.3.5 Disentangling Different OOD Generalization Behaviors . 192.4 Improving OOD Generalization of Neural Algorithmic Reasoning 212.4.1 Increase Dataset Size and Representativeness . . . . . . . 222.4.2 Improving Representation Distribution Shift . . . . . . . . 222.4.3 An Attention-Based Processor Competitive with MPNN-G 242.4.4 Hybrid Processor . . . . . . . . . . . . . . . . . . . . . . 252.5 Discussion on OOD Performances . . . . . . . . . . . . . . . . . 272.6 Related Works . . . . . . . . . . . . . . . . . . . . . . . . . . . . 282.7 Conclusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 293 Memorization Capacity of Multi-Head Attention in Transformers . 313.1 Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 313.2 Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . 333.3 Problem Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . 343.4 Main Results . . . . . . . . . . . . . . . . . . . . . . . . . . . . 363.4.1 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . 383.4.2 How tight is the lower bound on the rank? . . . . . . . . . 413.4.3 Analysis of the validity of assumptions . . . . . . . . . . 433.5 Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 453.5.1 Testing assumptions on Vision Transformer (ViT) . . . . . 453.5.2 Memory capacity increases with the number of heads . . . 463.6 Conclusion and Discussion . . . . . . . . . . . . . . . . . . . . . 494 Conclusion and Future Work . . . . . . . . . . . . . . . . . . . . . . 504.1 Conclusion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 50vii4.2 Future Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 504.2.1 Improving memorization bounds for Transformers . . . . 514.2.2 Closing the gap between theoretical and empirical memo-rization power of Transformers . . . . . . . . . . . . . . . 514.2.3 Establishing connection between expressivity and reasoning 514.2.4 Improving reasoning performance of models beyond archi-tectural modifications . . . . . . . . . . . . . . . . . . . . 52Bibliography . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 53A Appendix for OOD Generalization . . . . . . . . . . . . . . . . . . . 66A.1 Experiment Details . . . . . . . . . . . . . . . . . . . . . . . . . 66A.2 Leveraging Intermediate Steps of Algorithms . . . . . . . . . . . 67A.3 Encode-Process-Decode Framework . . . . . . . . . . . . . . . . 68A.4 Alternative Representation for Input Index Representation Shift . . 69A.5 Additional Approaches that Failed . . . . . . . . . . . . . . . . . 69B Appendix for Memorization . . . . . . . . . . . . . . . . . . . . . . . 78B.1 Proof of Lemma 1 . . . . . . . . . . . . . . . . . . . . . . . . . . 78B.1.1 Proof of Claims 1 and 2. . . . . . . . . . . . . . . . . . . 82B.2 Proof of Proposition 2 . . . . . . . . . . . . . . . . . . . . . . . . 84B.3 Extension of Theorem 1 to Attention with skip connection . . . . 84B.4 Additional Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . 85B.5 Upper bound on the rank of matrix Z with stronger assumptions . 88B.5.1 Synthetic experiments under general position assumption . 90B.6 Testing Assumptions on Natural Language Tasks . . . . . . . . . 90viiiList of TablesTable 2.1 Test scores of three processors (MPNN-FC, PGN, GAT) onthe CLRS benchmark proposed by [94] on all 30 algorithmsaveraged over three random seeds. The performances with hintsare lower on the majority of algorithms, for all processors. Themean is taken over all 30 algorithms, the wins are taken withineach processor. For each algorithm, a win is declared when thescore for that algorithm is higher than the alternative (i.e., withvs. without hints). . . . . . . . . . . . . . . . . . . . . . . . . 14Table 2.2 Test scores of MPNN-G processor on controlled datasets of9 graph algorithms, i.e., L-CLRS-Len, L-CLRS-Deg, and L-CLRS-Len-Deg. . . . . . . . . . . . . . . . . . . . . . . . . . 21Table 2.3 Summary of dataset configurations used in this paper. . . . . . 23Table 2.4 Performance of different positional encoding strategies on positions-order aware algorithms using MPNN-G processor on L-CLRSdataset. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24Table 2.5 Mean test scores of different processors averaged over all 24distinct algorithms on the L-CLRS dataset (performances of in-dividual algorithms are averaged over three random seeds). Thetwo rows correspond to using the scalar index and the randomscalar index for input representation respectively. The hybridmodels are a combination of MPNN-G and 2WL with the sameparameter budget as individual processors (Table A.7 and A.8contain the complete test scores of individual algorithms). . . . 26ixTable 2.6 Summary of sequentially applying the improvements proposed inthis work. Test scores are averaged over all 24 distinct algorithms. 26Table 3.1 Testing assumptions on ImageNet and ViT: while only one as-sumption is violated in the embedding layer, both hold true aftera single layer of Attention or Transformer. The General Positionassumption is violated in all cases. . . . . . . . . . . . . . . . 46Table A.1 Performance of different positional encoding strategies on positions-order aware algorithms for MPNN-FC on L-CLRS dataset. . . 70Table A.2 Full test results of Table 2.1. . . . . . . . . . . . . . . . . . . . 72Table A.3 Validation scores of Table 2.5. . . . . . . . . . . . . . . . . . . 73Table A.4 Validation scores of Table 2.2. . . . . . . . . . . . . . . . . . . 73Table A.5 Test scores of proposed processors on 12 representative algo-rithms on (low-data regime) CLRS dataset. While the 2WLprocessor shows some level of robustness to the low-data regime,MPNN outperforms 2WL in shortest-path-related tasks sincemax aggregation inductive bias is baked into the architecture,while the 2WL architecture requires a higher amount of data tolearn such inductive biases. However, the hybrid processor stillgeneralizes better compared to both processors. . . . . . . . . 74Table A.6 Test scores of different hybrid models for combinations of MPNN-G and 2WL on L-CLRS dataset. A hybrid model of two differentprocessors brings accuracy gain over each individual processor. 74Table A.7 Full test scores of Table 2.5 for Scalar Index. . . . . . . . . . . 75Table A.8 Full test scores of Table 2.5 for Random Scalar Index. TheWeighted Mean is average over all 30 algorithms when consider-ing repeated algorithms with higher weight. . . . . . . . . . . 76Table B.1 Important notations used in the paper. . . . . . . . . . . . . . . 78Table B.2 Testing assumptions on Wikipedia for BERT and GPT2. L0refers to the embedding layer, and L1 refers to the output of thefirst layer. . . . . . . . . . . . . . . . . . . . . . . . . . . . . 92xList of FiguresFigure 2.1 Encode-process-decode visualization for the BFS task. Theinput is a graph with the starting node index of two. The noderepresentations are encoded into a latent space. Then T consec-utive steps of processing are applied to the latent dimensionsusing a GNN. Finally, the output pointer for each node is de-cided by pairwise comparisons of node representations. . . . . 13Figure 2.2 (a) In-distribution (ID) accuracy (x-axis) vs OOD accuracy(y-axis) during different optimization steps for four different al-gorithms; in contrast to previous literature [65], either the linearcorrelation is non-existent, or the slope of the line is far frombeing one. (b) Failure of two widely-used data augmentationmethods when applied to neural algorithmic reasoning tasks.Top: dropping a single edge ([79]) on DFS (starting from node1) causes all of the output pointers in the graph (i.e. red arrows)to reverse direction. Bottom: Using Mixup [116] for sorting aconvex node-wise combination of two sequences leads to a newsequence to sort, rather than the node-wise convex combinationof labels of nodes of each sequence. . . . . . . . . . . . . . . 16xiFigure 2.3 Model interpolation on the Matrix Chain Order task. We trainthree models (named A, B, C) and pairwise interpolate theirweights to obtain ID and OOD accuracies. Each plot shows theinterpolation of a pair (e.g., when the coefficient is set to 0 and1 in (a), we use model B and A respectively; when it is in themiddle, we linearly interpolate weights from B and A). Whilemodels achieve similar ID accuracy, their interpolations showvery different OOD performances. . . . . . . . . . . . . . . . 19Figure 2.4 Illustration of two datasets generated for measuring the perfor-mance of a model on Bridges task. (a) The dataset generatedby the CLRS benchmark. (b) We generate a two-communitydataset for a stricter measurement of the Bridges task. Thedashed edge is present in one graph and absent in the othergraph of the pair. Red edges show the target bridges. Greenovals show connected components. A model with 98% OODF1 score on the dataset (a) fails to detect the bridge between thetwo communities of the dataset (b). . . . . . . . . . . . . . . 28Figure 3.1 Testing Kruskal Rank of query tokens on the output of one layerRandom Attention on ImageNet dataset. The Kruskal Rank isonly slightly larger than n (Assumption 1), and much smallerthan d (General Position). . . . . . . . . . . . . . . . . . . . 46Figure 3.2 Testing memorization as a function of heads (a, b) and contextsize (c, d) for both classification and regression, under Assump-tions 2 and 1. Synthetic examples with d = 64 and sharedcontext (see Proposition 1). Both settings show a positive corre-lation between heads/context size with memorization capacity,while the classification setting shows a linear relationship withboth. The variance in regression plots is magnified in somecases due to the log scale of the plot. . . . . . . . . . . . . . . 48xiiFigure A.1 Model interpolation charts on three additional tasks. For eachtask, we train three models and pairwise interpolate their weightsand plot ID and OOD accuracies. Each plot shows the inter-polation of one pair for each task. As suggested by the plots,ID and OOD accuracies have very different landscapes and IDaccuracy is a suboptimal model selector. In (a), the right plotshows the highest OOD accuracy corresponds with the lowestID accuracy. In (b), the middle plot achieves its highest OODaccuracy when ID accuracy is not at its peak. In (c), the rightplot shows different rates of change for ID and OOD accuracies,OOD accuracy is at its peak when ID is not, and for coefficient0.5, the OOD accuracy is higher than that of ID. . . . . . . . 77Figure B.1 Testing memorization under general position assumption forboth context and query vectors. While the linearity in H stillremains, the linearity in n no longer strictly holds, suggestingH(d−1)+1 to be a better alternative under general positionassumption. . . . . . . . . . . . . . . . . . . . . . . . . . . . 91xiiiAcknowledgmentsI am deeply indebted to my esteemed supervisors, Renjie Liao and Christos Thram-poulidis, for their remarkable patience, unwavering dedication, and invaluableguidance throughout this research endeavor. Their expertise, passion for knowledge,and profound insights have been a constant source of inspiration. Their mentorshiphas shaped not only my academic growth but also my personal development.My gratitude extends to my exceptional collaborators, Kevin Swersky, ThomasKipf, and Milad Hashemi. Working alongside these brilliant minds has been anenriching experience. Kevin’s boundless creativity and ability to generate freshideas pushed the boundaries of our research, while Thomas’s remarkable intuitionand insightful perspectives consistently offered new directions for exploration.Last but certainly not least, I wish to extend my deepest thanks to my parents.Their unwavering support, love, and understanding have been the bedrock uponwhich my academic pursuits have thrived. Their sacrifices and willingness to puttheir children’s dreams and aspirations before their own have been a constant sourceof motivation and strength for me.xivChapter 1IntroductionDeep Learning has emerged as a highly successful approach for solving a wide rangeof problems in various domains, including computer vision [24], healthcare [81],and natural language processing [17]. At the core of this success is the Transformerarchitecture [89], which has become the preferred choice for achieving state-of-the-art results in these fields [29, 76]. The recent advancements in scaling Transformersand the availability of extensive data have paved the way for models like GPT-4,which have achieved near human-level performance on challenging tasks such asmedical and law exams [70]. These remarkable achievements by large-scale modelshave even sparked discussions about their potential to exhibit characteristics ofArtificial General Intelligence (AGI) [19].One of the key signs of artificial general intelligence is to show the ability toreason [37]. There have been benchmarks proposed to measure various forms of rea-soning in neural networks, e.g., visual reasoning [47], commonsense reasoning [83],algorithmic reasoning [94]. Here, we consider algorithmic reasoning since it has aclear well-defined definition as opposed to most other types of reasoning.One key property of large language models (LLMs) is that they have beentrained on massive amounts of data from the internet. As a result, it is unclearwhether their high performance on tasks results from seeing evaluation data pointsin their pre-training data (and hence memorization) or reasoning ability. Recentworks have cast doubt on the general reasoning capability of models such as GPT-4.Razeghi et al. [78] show that the performance of LLMs in arithmetic has a positive1correlation with the number of times LLMs have seen a number in their trainingdata. Wu et al. [105] show state-of-the-art LLMs lose significant performance ontasks slightly modified to depart from the pre-training data but are algorithmicallythe same (e.g. non-base-10 arithmetic). Finally, it has been observed that GPT-4’s performance in solving competitive programming questions is significantlylower when tested on problems that were posted after GPT-4’s training period [41].In summary, while LLMs achieve human-level performance on many tasks dueto their vast pre-training, their reasoning ability remains limited. Performancestrongly depends on the overlap between pre-training and evaluation data and out-of-distribution generalization and reasoning remain open challenges.In this thesis, our contributions are as follows:• In Chapter 2, we investigate whether Graph Neural Networks and Transform-ers have the ability to learn algorithmic reasoning in a data-driven manner. Weidentify the challenges and investigate the out-of-distribution generalization(OOD) in Neural Algorithmic Reasoning tasks, showcasing key differencesfrom common OOD tasks. Specifically, we find that while models can achieveperfect in-distribution accuracy, out-of-distribution performance is highly un-predictable and poor. Furthermore, we propose a series of improvementsincluding a Transformer-based architecture to enhance the performance onthe benchmark at hand. Finally, we show that while our improvements signifi-cantly boost performance on the benchmark, the core question of algorithmicreasoning remains unsolved. Simple changes to the OOD test set can de-crease model performance from near-perfect to less-than-chance accuracy.This shows that despite a high amount of expressivity, they fall short of trulyeliciting reasoning abilities.• The recent findings on memorization vs. generalization, together with resultsin Chapter 2 point that recent success in increased reasoning may be largelydue to the high capacity of models to retain information. This motivates us totake a step back and examine the expressivity and capacity of Transformers,and ask “To what extent is a Transformer capable of memorizing data?”. Toshed light on this question, in Chapter 3 we theoretically analyze the abilityof Transformers to memorize arbitrary data points. Since attention is central2to Transformers, we consider a softmax-attention layer and provide a lowerbound on multi-head attention memorization capacity in Transformers. Wethen test the assumptions we made and show they often hold in practice.Finally, to validate our theoretical findings, we perform synthetic experimentsand show a linear relationship between memorization capacity and the numberof attention heads.The remainder of this chapter introduces concepts necessary for subsequent chapters.1.1 BackgroundIn the modern era of data-driven applications, a vast amount of information isavailable in diverse formats, often presenting complex structures. These dataformats range from molecular representations in bioinformatics to natural languageand social networks. Effectively extracting information and patterns from suchstructured data has been a paramount challenge in machine learning.Early attempts at processing structured data predominantly relied on RecurrentNeural Networks (RNNs) [43], which demonstrated notable success in varioustasks, such as machine translation [104] and early algorithmic reasoning tasks [120].RNNs proved particularly adept at handling sequential data, thanks to their abilityto capture temporal dependencies and context. Consequently, they emerged asa leading choice for processing sequential data and transforming it into usefulrepresentations. However, despite their achievements, RNNs came with their ownset of challenges such as being notoriously hard to train and scale.Recently, Graph Neural Networks (GNNs) [54] and Transformers [89] haveemerged as compelling alternatives for processing structured data, significantly im-pacting the landscape of machine learning. These models have garnered substantialattention and have become the preferred solutions for a diverse array of tasks. Theirhigh performance has even extended their usability beyond graphs/sequences and toimage data, exemplified by their prominent utilization as state-of-the-art (SOTA)models in image classification tasks on benchmarks like ImageNet [24, 29]. Inwhat follows, we formally define Transformers and Graphs Neural Networks, theconnection between these two, and how they are used in relevant applications.31.1.1 Graph Neural NetworksGraph Neural Networks (GNNs) have emerged as effective deep learning models toprocess graph data. Given an input graph G with n nodes and d-dimensional nodefeatures X ∈ Rn×d , and adjacency matrix A ∈ Rn×n, the goal of these models is toprocess the graph and produce useful representations for various purposes such asclassification and generation. A GNN layer can be as simple as follows:GNNW (X ,A) = AXW ,where W ∈Rd×d1 is a learnable matrix. At the output of this model, each node has ad1 dimensional output, which is the average of its neighbors in the graph, multipliedby matrix W . More broadly, message-passing GNNs have the following genericform:m(i) = AGGR({X j| j ∈N (i)})(1.1)z(i) = UPDATE(X i,m(i)) , for all i = 1,2, . . . ,n (1.2)where N (i) is the set of neighbors of node i, AGGR is an aggregation function,UPDATE is a shared update function, and z(i) is the output of the GNN for node i.In the case of our simple GNN, AX acts as the aggregation function, and producingthe output representation by multiplying by W is the update function.1.1.2 TransformersThe Transformer architecture [89] was initially developed to process sequencesof text data for applications such as machine translation. However, later, thisarchitecture proved to be highly effective in processing other modalities such asimages [29], audio [76], and algorithmic data [107].Interestingly, it has been shown that Transformers fall within the generic formof message-passing GNNs defined in Eq. (1.1) and Eq. (1.2) [48]. Formally, aTransformer with a single Attention head takes input matrix X ∈Rn×d and produces4outputs as follows:H = Softmax(XW Q(XW K)⊤√d)X︸ ︷︷ ︸AGGRW V (1.3)X ′ = LN(X +H), (1.4)Z = LN(X ′+MLP(X ′)), (1.5)where W Q,W K ,W V ∈ Rd×d are learnable matrices, LN is the Layer Normalizationoperation, Softmax is the softmax function applied row-wise, MLP is a Multi-LayerPerceptron, and Z ∈ Rn×d is the output matrix.This shows that the Transformer can be viewed as a GNN with feature nodes X ,fully connected adjacency matrix, aggregation function defined in Eq. (1.3), and therest of the operations as the update function.1.1.3 Neural Algorithmic ReasoningNeural Algorithmic Reasoning is the task of mimicking classical algorithms usingneural networks. In such tasks, given a ground-truth algorithmic function f :X →Y , the goal is to train a neural network to mimic this ground-truth function.One distinction between algorithmic tasks and natural tasks is that in algorithmictasks typically data generation and function evaluation is cheap. This is while innatural tasks both of these require human intervention (e.g., image classificationrequires both data collection and labeling). As a result, algorithmic tasks can benefitfrom infinite data generation and have the flexibility to define tasks. However,the evaluation metric is also stricter, as the trained neural networks are typicallyrequired to show strong generalization and generalize beyond the distribution theyhave been trained on.1.1.4 Large Language ModelsTraditionally, SOTA models in natural language processing utilized supervisedlearning on task-specific data to achieve high performance. Radford and Narasimhan[74] showed first pretraining a Transformer in an unsupervised manner, then fine-tuning it on the underlying tasks results in superior performance. Later, Radford et al.5[75] showed that by training larger Transformers, a language model can achieve highperformance even without fine-tuning on the underlying tasks. Recently, it has beenshown that with further scaling of data and computing such an unsupervised pre-training on a large dataset can achieve SOTA results without any fine-tuning [17, 70].Formally, given a dataset D of text documents, the training objective of aTransformer language model TFW with parameters W is the following:argminWE(x1,x2,...,xn)∼P(D)[n∑i=2ℓ(TFW (xi |x1,x2, . . . ,xi−1))],where P(D) is a uniform distribution over contiguous windows of n (sub)words fromthe documents within D . TFW (xi |x1,x2, . . . ,xi−1) represents the output probabilityof the Transformer predicting the next word given the input sequence x1,x2, . . . ,xi−1.The loss function ℓ is typically the cross-entropy loss, aiming to maximize theprobability of predicting the correct next word xi.Remarkably, large models like GPT-3 have demonstrated that this trainingobjective, also known as the next-word-prediction objective, coupled with substantialamounts of data (|D |) and model parameters (|W |), becomes remarkably powerful,obviating the need for fine-tuning to achieve state-of-the-art performance, oftensurpassing even human-level performance in various tasks.This paradigm shift in NLP from supervised objectives to unsupervised objec-tives has raised intriguing questions, such as how these models develop task-solvingbehaviors solely from an unsupervised learning objective.1.2 Related Work1.2.1 Out-of-Distribution Generalization and ReasoningThere has been a long quest for going beyond the distribution the models have beentrained on. A line of work has focused on the robustness of neural networks on imagedata when tested on data points conceptually the same as the training set, but driftingfrom the training set in some form, such as temporal or even adversarial [61, 101].These works typically find various forms of data augmentation to be highly effective,especially in distribution shifts that involve corrupting images [42, 116].6Another line of research has focused on solving algorithmic problems usingneural networks. These works range from solving simple arithmetics [56, 118], tothe ambitious goal of making general-purpose computers such as a Neural TuringMachine [38]. In Section 2.3.1, we provide a more thorough body of related workon reasoning.1.2.2 In-Distribution GeneralizationIn-distribution generalization is typically what is considered in the majority ofmachine learning works. This type of generalization assumes that the evaluationset is sampled iid from the same distribution as the training set. This then allowsfor performance guarantees given enough samples. For instance, there have beennumerous generalization bounds proposed for neural networks [6, 61, 69]. Modelslike large language models, however, are typically evaluated with non-iid test data,as a result, it is not easy to define in-distribution generalization, or any generalizationbounds as they violate this assumption.1.2.3 MemorizationOn the empirical side, there has been a line of work focusing on privacy, andgeneralization. Carlini et al. [20] showed an attack to extract training data from atrained GPT-2 model. For instance, they showed that it is possible to extract theexact MIT public licenses from this model and it has completely memorized theselicenses. Furthermore, Zhang et al. [114] showed that modern neural networksare highly overparameterized to the extent that they can fit random labels from thetraining data.On the theoretical side, the memorization of fully-connected networks (espe-cially ReLU networks), has been extensively analyzed. Baum [8] showed that asingle-hidden-layer Threshold-activated neural network can memorize N exampleswith binary labels with O(N) parameters. Bubeck et al. [18] showed a similarcomplexity of O(N) for ReLU-activated networks with real-valued labels. Laterworks extended the upper-bound for the number of parameters for multi-layer ReLUactivated networks [95, 110]. Sontag [82] proved a lower bound on the necessarynumber of parameters needed to shatter a dense class of examples using piece-7wise analytic definable functions (which includes ReLU). This work showed thatmemorization for a large class of examples is only possible with at least Ω(N)parameters. Harvey et al. [40] proved a nearly tight (up to a log factor) bound ofΘ˜(WL) on the VC dimension of ReLU networks with W parameters and L layers.This work gives a lower bound of Ω˜(N) parameters for a fixed network depth (L),and Ω˜(√N) parameters for a Θ˜(√N)-deep network. Subsequent works aimed atclosing the gap between this sub-linear (Θ˜(√N)) parameter-count, and the linear(Θ˜(N)) parameter-count upper-bound on the memorization. Park et al. [72] provedthe first sublinear parameter count with O(N2/3) parameters by adding a norm-basedseparability assumption, which bypasses the results of Sontag [82]. Later, Vardiet al. [87] closed the gap and proved O˜(√N) parameters are sufficient to memorizeany set of examples satisfying norm-based separability. Interestingly, this shows thatmemorizing a large class of N-examples is not harder than shattering a particularset of N data points.8Chapter 2Out-of-DistributionGeneralization for NeuralAlgorithmic Reasoning2.1 IntroductionAlgorithms have been a vital part of today’s computing systems, solving tasksranging from scheduling to finding the shortest path in real-world maps. Algorithmsoften come with benefits like performance guarantees, provable correctness, andgeneralization to various instances. Designing algorithms to solve tasks, however,can be quite hard in many situations. For instance, given any task, a computerscientist has to either follow some existing algorithm paradigms (e.g., divide andconquer) or come up with some new algorithm, which typically requires significanthand-design efforts. Moreover, some problems lie in the NP-hard class of whichthe solutions require an exponential amount of time for execution, and designingefficient approximate algorithms for such problems often requires several years ofresearch. Finally, executing an algorithm often requires preprocessing data intoa predefined format before feeding the data into a computer program. On theother hand, deep learning has been successful in solving tasks of many domainswithout much task-specific adaptation. For instance, a Transformer [89] could solve9image classification, machine translation, and language modeling with negligiblechanges to its architecture. Moreover, deep learning models are well known for theirability in solving tasks in an end-to-end fashion, thus removing the necessity of datapreprocessing and feature design. Neural networks, however, are not resilient todistribution shifts, and their symbolic manipulation and reasoning capabilities arelimited [115].Therefore, it is natural to ask whether one can leverage deep neural networksto mimic algorithms, i.e., neural networks execute the same computation as givenalgorithms. The hope is to combine the strengths of both worlds, e.g., 1) performanceguarantees, generalization to various instances, and transferability across tasks fromalgorithms, and 2) end-to-end training, universal approximation, and fast parallelcomputation from neural networks. Recently, researchers have been working on thisproblem and the area is dubbed as Neural Algorithmic Reasoning [90]. Here, thegoal is not only to enhance the symbolic manipulation and reasoning capabilities ofneural networks but also to leverage them to find more efficient and faster algorithms.Existing works have shown that this problem is challenging, especially whenthe model is evaluated on test samples that are drawn from a different distributionthan train samples (i.e., it is hard to generalize to out-of-distribution data) [94].For example, while deep neural networks could be trained to perform sorting onsequences of a given, fixed length, when faced with longer sequences, their perfor-mance drops substantially [88]. In this work, we analyze the out-of-distribution(OOD) generalization of neural networks on algorithmic reasoning tasks. We firstinvestigate several challenges with OOD generalization in our problem context andpropose solutions to some of the issues that arise. In particular, our contributionsare as follows:• In Section 2.3, based on the CLRS benchmark, we show how OOD set-tings of neural algorithmic reasoning tasks are different from common ones(Section 2.3.1) and identify several key challenges: input distribution shift(Section 2.3.2), unrepresentative dataset generation in both train and test(Section 2.3.3), model selection difficulties due to near-perfect in-distributionvalidation metrics (Section 2.3.4), and a systematic investigation of OODgeneralization under controlled distribution shift of graph sizes and node10degrees (Section 2.3.5).• In Section 2.4, we propose solutions to some of these challenges, including1) enlarging the in-distribution training dataset to push the limits of OODperformance (Section 2.4.1), 2) a simple-yet-effective fix to input distribu-tion shift (Section 2.4.2), 3) an attention-based 2-WL graph neural network(GNN) [67] processor that significantly outperforms the state-of-the-art whencombined with message passing GNNs (Sections 2.4.3 and 2.4.4 ).2.2 BackgroundWe follow the setup of the CLRS benchmark [94]. All algorithmic tasks aretranslated and represented as graphs (possibly empty or fully-connected) with task-specific node/edge/graph feature inputs and node/edge/graph-level target outputs.Each task also contains the ground truth values of all variables presented in interme-diate steps of an algorithm, which are called hints and could be used as additionalsupervision to train the model. The aim of hints is to help mimic individual steps ofthe underlying algorithm and finally predict the target output.Each instance of a task contains a graph G = (V ,E ) with a set of n nodesV = {v1,v2, . . . ,vn} and m edges E = {ei = (vi1 ,vi2) : vi1 ,vi2 ∈ V ,1≤ i≤ m} withan adjacency matrix A ∈Rn×n, edge feature tensor X E ∈Rn×n×dE , and node featurematrix X V ∈ Rn×dV . We eschew an explicit graph feature vector as it could berepresented as a concatenation of features to the node feature matrix. For a node(column) vector zv, we use capital from Z to indicate a vertically stacked matrix(one row per node) of all node vectors across the graph. Moreover, we denotecolumn-wise concatenation of two matrices/row-vectors E and F as [ E , F ].All models we consider are GNNs with an encode-process-decode framework[7]. For an algorithm with T iterative steps, we consider a recurrent GNN whosemessage-passing steps correspond to algorithm steps. At each step t (1≤ t ≤ T ) ofprocessing, all input features are embedded into embedding vectors zinpv ∈ Rdh bytheir corresponding encoders, for each individual node v ∈ V . The processor thentakes the embedding vectors and the hidden states at step t and outputs the hiddenrepresentations of step t+1. Formally, a (simplified) formula for the processor is as11follows:H(t+1) = GNN([ H(t), Z inp ],A), (2.1)where GNN(X ,A) represents the one-step computation process of a GNN processorwith input node features X and adjacency matrix A. Finally, outputs are decodedusing a simple MLP per node/edge, with shared parameters. The loss value isevaluated based on the outputs of decoders at iteration T .To better illustrate this framework, we explain how Breadth-first Search (BFS) isformulated in this framework. Given an input undirected graph and a starting node,the task is to predict for each node the pointer to its predecessor when applyingBFS on this graph. Namely, we have two-dimensional node features X V ∈ Rn×2where the first dimension determines the index of a node and the second dimensiondetermines whether the node is the starting node. We also have the adjacency matrixA of the graph. The encoder is a linear layer with a weight matrix W enc ∈ R2×dhwhere dh is the latent embedding size. Then, for the processing step, T recurrentsteps of GNN are applied to the latent space with adjacency matrix A to get thefinal latent node representations H(T ) ∈ Rn×dh . Finally, for each pair of nodesi, j (1≤ i, j ≤ n), a score αi, j determines how likely i points to j as the predecessorand yˆi = argmax jαi, j is the predicted output pointer of node i. To compute αi, wecan apply an MLP to [h(t)i ,h(t)j ] followed by a (row-wise) Softmax non-linearity.The loss for this task is a cross-entropy loss between predicted row scores αi,: andthe true (one-hot row-wise) pointer scores for each node i. This process is illustratedin Figure 2.1. For additional examples, see Appendix A.3.As the processor is considered the central part of this architecture, here weexplain the three processor baselines used by the benchmark:• MPNN [36]: A message passing neural network with a max-aggregation. Onthe benchmark, this processor assumes a fully connected graph regardless ofwhether the downstream task contains any graph or not. We regard this fully-connected MPNN processor as MPNN-FC. We also consider the originalMPNN which operates on given graphs (if graphs are unavailable for a task,a fully connected graph is used), and call it MPNN-G.1221 34521521 34521 34521 345Enc Proc Dec1 21 4PointerNoPointerNode EmbeddingEncoderDecoderMessage PassingFigure 2.1: Encode-process-decode visualization for the BFS task. The inputis a graph with the starting node index of two. The node representationsare encoded into a latent space. Then T consecutive steps of processingare applied to the latent dimensions using a GNN. Finally, the outputpointer for each node is decided by pairwise comparisons of node repre-sentations.• PGN [92]: Similar to MPNN, but PGN operates on a graph derived fromboth the input and hints. When no hints are available, and the task does notcontain any graph, it assumes an empty graph and hence behaves the same asDeepSets [113].• GAT [91]: Similar to PGN, but GAT replaces the max-aggregation-basedmessage passing in MPNN with the attention based one.To fairly benchmark these processors, the encoder and the decoder are controlled tobe the same.2.3 Current Challenges of OOD generalization of NeuralAlgorithmic Reasoning TasksIn the CLRS, Velicˇkovic´ et al. [94] measures both the in-distribution (ID) and theout-of-distribution (OOD) generalization of the aforementioned processors over30 classical algorithms. For example, they generate validation graphs from thesame distribution as training but generate testing graphs from different distributions,e.g., larger sizes in terms of the number of nodes. Table 2.1 contains the OOD testaccuracies of the algorithms when using hints (i.e. intermediate steps). We also13Table 2.1: Test scores of three processors (MPNN-FC, PGN, GAT) on theCLRS benchmark proposed by [94] on all 30 algorithms averaged overthree random seeds. The performances with hints are lower on the majorityof algorithms, for all processors. The mean is taken over all 30 algorithms,the wins are taken within each processor. For each algorithm, a win isdeclared when the score for that algorithm is higher than the alternative(i.e., with vs. without hints).Method MPNN-FC PGN GATUse Hints Yes No Yes No Yes NoMean 51.02 60.04 52.31 57.02 44.69 49.53Wins 8/30 22/30 10/30 20/30 14/30 16/30include the performance of the different models when disabling hints. As the tableshows, in the current form of benchmark, hints do not bring any performance gainon the final output for the majority of algorithms. Moreover, to learn algorithmsbeyond those in the CLRS benchmark, hints might be unavailable or infeasible tocollect at all (e.g., NP-hard problems). Based on these considerations, we deactivatehints for the rest of the paper and instead focus on end-to-end training with onlythe supervision of the final output. In the following, we highlight the fundamentalchallenges in the OOD generalization of neural algorithmic reasoning in the CLRS.2.3.1 Difference with Common OOD SettingsSeveral recent works have studied the problem of OOD generalization in both imageand graph domains [4, 5, 21, 71]. However, the nature of OOD generalizationin neural algorithmic reasoning is fundamentally different from these works. Toillustrate an example, consider the phenomena of accuracy on the line [65] forimage classifications where some works have observed the OOD performance hasa linear relationship with (and often is close to) ID performance in some imageclassification benchmarks. However, in our setting, we often observe perfect IDgeneralization with poor (sometimes even less than random guesses) OOD accuracy,and OOD generalization might keep improving after the ID validation performancecompletely saturates (see Bridges task in Figure 2.2a). Moreover, popular data14augmentation techniques as an effective remedy for OOD generalization [39, 101],are generally inapplicable to our setting. For instance, Mixup [116] assumes convexcombinations of inputs have the same convex combination of outputs, but for manyalgorithms, a convex combination of inputs from two tasks creates a new task withoutput significantly different from the convex combination of outputs from two tasks(e.g., sorting the average of two sequences creates a whole new task, and the outputis not a linear function of outputs of two tasks), thus violating the assumptions ofMixup. Even graph-specific data augmentation methods are generally inapplicable.For instance, DropEdge [79] assumes by randomly dropping an edge in the graph,the target output stays the same. However, for many graph algorithms, altering anyedge could change the output of the underlying task significantly. For example,performing DFS on a ring graph is an extreme failure example, since removing asingle edge changes the outputs (e.g., pointers of nodes to be visited next) of allnodes in the graph. See Figure 2.2b for an illustration of the failure modes of suchexamples.A recent line of work has particularly investigated OOD generalization inGNN settings. However, the assumptions they make about generalization acrossenvironments (in our case graph sizes) are invalid to the neural algorithmic reasoning.Yehudai et al. [109] studied such generalization from a local k-hop neighborhoodperspective (where k is the number of GNN layers) and connected failures ingeneralization to distribution shifts in the such local neighborhood. For algorithmictasks, the number of reasoning steps are so large (sometimes larger than the numberof nodes) that such a local neighborhood often contains the whole graph andthe analysis would be vacuous. Some prior works have also analyzed the OODgeneralization through the lens of graphons [10, 64, 119] where graphs alongwith labels are drawn from some grahpon models. In these settings, the implicitassumption is that if we restrict the original graph to a subgraph induced by asubset of its nodes, the labels within the subgraph stay the same. However, thisfavorable property does not hold in neural algorithmic reasoning tasks. For instance,consider the DFS example in Figure 2.2b. By removing node “2” and restrictingthe graph to the rest of the nodes, all the pointers turn into a clockwise order; hencethe outputs change when restricting a graph to a smaller subset of its nodes. Inalgorithmic reasoning tasks, the relation between ID and OOD generalization is150 20 40 60 80 100020406080100OOD AccuracyBridgesLinear Fit(Slope: 0.89, R2: 0.31)0 20 40 60 80 100020406080100 DfsLinear Fit(Slope: 0.26, R2: 0.85)0 20 40 60 80 100ID Accuracy020406080100OOD AccuracyMst PrimLinear Fit(Slope: 1.06, R2: 0.85)0 20 40 60 80 100ID Accuracy020406080100 Quick SortLinear Fit(Slope: 0.11, R2: 0.53)(a) ID vs OOD accuracy relation1234561234560.4 0.5 0.60.6 0.3 0.00.5 0.4 0.3DropEdgeMixup(b) Data augmentation failure modesFigure 2.2: (a) In-distribution (ID) accuracy (x-axis) vs OOD accuracy (y-axis) during different optimization steps for four different algorithms;in contrast to previous literature [65], either the linear correlation isnon-existent, or the slope of the line is far from being one. (b) Failureof two widely-used data augmentation methods when applied to neuralalgorithmic reasoning tasks. Top: dropping a single edge ([79]) on DFS(starting from node 1) causes all of the output pointers in the graph(i.e. red arrows) to reverse direction. Bottom: Using Mixup [116] forsorting a convex node-wise combination of two sequences leads to a newsequence to sort, rather than the node-wise convex combination of labelsof nodes of each sequence.much more intricate than the above assumptions; the only connecting point betweenthe target label of each task across different graph sizes is they are generated usingthe same algorithm. This makes even formulating the OOD generalization and thedesired invariance properties in such tasks nontrivial. Although some prior workshave discovered some desired invariance properties of a model for stronger OODgeneralization of some algorithmic tasks, such as homogeneity of the neural network[84], finding sufficient conditions remains an open question.2.3.2 Distribution Shift in Input RepresentationsBeyond improving models, finding suitable data representations to facilitate OODgeneralization is also crucial. For instance, all algorithms in the CLRS benchmarkreceive the node indices as a part of their input. For several tasks, this node index16serves as a unique flag (e.g., Convex Hull, Bridges), for others, the ordering playsan important role in the algorithm behavior (e.g., DFS algorithm starts with a nodewith the lowest index, then, at each step, an unexplored neighbor with the lowestindex gets expanded). As a result, the model must be able to compare node indicesto be able to execute the algorithm.The current node index implemented in the CLRS [94] appends a scalar posvi =(i−1)/n to the feature vector of each node vi with index i,(1≤ i≤ n). This kindof representation makes it harder for a processor to generalize when such orderingis important in a task. For instance, despite the model being able to reach perfectID accuracy to execute DFS on 16-node graphs; it fails to do so when we feed theinduced subgraph of the first quarter nodes of a 64-node graph. This is while thegraphs are still in distribution, and the only difference is that the indices have shiftedfrom 16 equispaced indices in the range [0,1] to 16 equispaced indices in the range[0,0.25]. In Section 2.4.2 we explore two different possible mitigations to such adistribution shift for node indices.2.3.3 Unrepresentative Dataset GenerationThe graphs in the CLRS benchmark are sampled using Erdo˝s–Rényi (ER) modelwith a fixed probability p (where p = 0.5 in most cases) and node/edge/graph-levelscalar input features are generated uniformly at random (mostly from [0, 1]). Inthis strategy, sampling is uniform and all possible graphs (of the same size) have anequal probability of being present in the dataset. However, it takes many samplesfor this random graph model to sample diverse graphs (e.g., node-degree and thenumber of edges in ERs concentrate around (n−1)p and n(n−1)p/2 respectively).This raises two issues, especially in the low-data regime (i.e., fewer training data).First, the training data might not be enough to uniquely identify the under-lying algorithm (e.g., BFS and DFS have the same output pointers on the classof tree graphs – all nodes point to their parents towards the root) and the OODgeneralization should not be expected in this case.Second, and perhaps more importantly, the test dataset may turn into a poorperformance measure as it is not diverse enough and biased. This bias may allowcheating solutions to gain higher scores, especially when extrapolating [3]. For17instance, two of the most challenging algorithmic tasks in the CLRS are ArticulationPoints and Bridges. In these tasks, the goal is to identify the nodes/edges whoseremoval increases the number of connected components of the underlying graph.Here, if a graph is too sparse, trivial bridges/articulation points appear in the graph.On the other hand, a dense graph leads to no such nodes/edges. In Section 2.4.4,we report a 98.69% OOD accuracy on the Bridges task using a hybrid combinationof 2WL and MPNN-G processor. At first sight, the Bridges algorithm seems to becompletely solved and the model can reliably find the bridges in a graph. However,the test dataset is not rich enough for testing the model. The bridges in the datasetare mostly chain subgraphs branching out of a connected component (see Section2.5 for more discussion).2.3.4 Model Selection is Hard with Nearly Perfect ID PerformancesGenerally, model selection in OOD settings is not as straightforward as in the IDsetting. This is because the validation data is not identically distributed as the testdata. This problem is accentuated even more in algorithmic reasoning tasks. Whentraining a model on algorithmic tasks of the CLRS benchmark, it usually fits theID validation data almost perfectly in terms of prediction accuracy. This leaves aharder model selection question than common OOD settings, i.e., “How to pick thebest model among all these perfect models?”. These perfect-ID models behave verydifferently on OOD data. For example, despite all the architectures achieving 100%ID accuracy on the Articulation Points task (Table A.3 in the Appendix), the OODaccuracies widely differ across different architectures, ranging from 69% to 88%.This issue exists even when comparing models from the same architecture but indifferent optimization steps. For example, for some tasks, the OOD accuracy keepsimproving even after perfectly fitting the ID data. An extreme case can be seen inFigure 2.2a for the Bridges algorithm. Here, the OOD accuracy keeps improvingfrom around 30% to above 90% after an almost perfect ID performance is achieved.Furthermore, OOD performances in some algorithms have a high varianceacross models, and sometimes even negatively correlate with ID performances.Inspired by recent works on mode-connectivity and model weight interpolation[31, 103], we interpolate between the weights of models to better understand the180.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0AccuracyIDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0IDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0IDOOD(a) Model B to A.0 .1 .2 .3 .4 .5 .6 .7 .8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0AccuracyIDOOD.0 .1 .2 .3 .4 .5 .6 .7 .8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0IDOOD.0 .1 .2 .3 .4 .5 .6 .7 .8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0IDOOD(b) Model C to A0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 .8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0AccuracyIDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 .8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0IDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 .8 0.9 1.0Interpolation Coefficient0.40.50.60.70.80.91.0IDOOD(c) Model C to BFigure 2.3: Model interpolation on the Matrix Chain Order task. We trainthree models (named A, B, C) and pairwise interpolate their weights toobtain ID and OOD accuracies. Each plot shows the interpolation ofa pair (e.g., when the coefficient is set to 0 and 1 in (a), we use modelB and A respectively; when it is in the middle, we linearly interpolateweights from B and A). While models achieve similar ID accuracy, theirinterpolations show very different OOD performances.loss landscape of such tasks. We pick the Matrix Chain Order task which shows ahigh variance even when the initialization point of the model is fixed. Here, despiteall models attaining 99.5% ID accuracy, some of them have an OOD accuracy ofaround 60%, and some have an OOD accuracy of around 80%. Figure 2.3 showspairwise interpolation of three models trained on this task. According to the figure,there are barriers when connecting two sets of weights. Quite counter-intuitively,the ID barrier between two models with a larger OOD gap is the smallest (Figure2.3b). Moreover, there are two barriers for the ID performance and one barrier forthe OOD in Figure 2.3c. Finally, when interpolating from model C to model A, theID and OOD do not show any positive correlation and even negative correlationin the first half of the plot (2.3b). These all support the fact that the ID validationperformance is not a suitable measure for model selection, and ID and OOD dataclearly have different loss landscapes. See Figure A.1 in the Appendix for additionalinterpolation plots.2.3.5 Disentangling Different OOD Generalization BehaviorsThe CLRS benchmark [94] generates the underlying graphs using an Erdo˝s–Rényimodel with a fixed probability p. This means that the node degree distribution19increases as the graph size increases. While prior works have worked on provingOOD generalization guarantees when the local structure of graphs stays the same[109], neural algorithmic reasoning is different from standard graph benchmarks.Here, when the graph size increases, higher-depth reasoning is needed. As a result,a fixed-depth local structure of the graph is no longer sufficient for a node to decidethe label. To disentangle these two types of generalization, we create three datasetconfigurations of the CLRS benchmark to better understand such generalizationbottlenecks.We consider only graph algorithms since it is not possible to control the message-passing locality of sequence/set algorithms. There are 11 distinct graph algorithmsin the benchmark. We remove the Articulation Points and Bridges algorithmssince they require complicated sampling/evaluation strategies to prevent cheatingsolutions from succeeding. We create the following datasets:• L-CLRS-Len: We control the degree distribution of nodes between trainingand testing to be the same via sampling random K-regular graphs. Only thesize (i.e., number of nodes) of testing graphs is larger than those in training.This is an ideal setting for studying size generalization.• L-CLRS-Deg: We control the number of nodes between training and testingto be the same. Only the degree distribution of nodes changes, which isachieved by changing the K value in K-regular graphs. This is ideal forstudying how the degree distribution shift affects the generalization.• L-CLRS-Len-Deg: We allow both the size and the degree distribution to bechanged between training and testing, which is a setting closest to the originalCLRS. It is useful to study both size and degree generalization. Here, in thetest set, we not only increase the number of nodes but also increase the Kvalue of K-regular graphs.The details of these datasets are summarized in Table 2.3. Note that the mo-tivation behind choosing K-regular graphs is that degree shift can be completelycontrolled and no shift in the variance of degree. Moreover, dataset leakage isavoided as we generate more training graphs (i.e., graphs in the test dataset havezero probability of appearing in the training set). In addition to the CLRS’s own20Table 2.2: Test scores of MPNN-G processor on controlled datasets of 9 graphalgorithms, i.e., L-CLRS-Len, L-CLRS-Deg, and L-CLRS-Len-Deg.Accuracy Metric Node Level Graph LevelMethod CLRS-Len-Deg CLRS-Len CLRS-Deg CLRS-Len-Deg CLRS-Len CLRS-DegBellman Ford 99.48± 0.07 99.57± 0.06 99.40± 0.09 84.93± 2.24 87.53± 1.69 82.97± 2.11Bfs 100.00± 0.00 99.92± 0.04 100.00± 0.00 100.00± 0.00 97.77± 1.23 100.00± 0.00Dag Shortest Paths 99.84± 0.04 99.94± 0.00 99.81± 0.02 95.10± 1.18 98.17± 0.12 94.37± 0.66Dfs 74.63± 3.91 81.76± 3.66 98.91± 0.23 0.13± 0.12 3.67± 1.68 81.57± 2.82Floyd Warshall 63.84± 1.07 64.51± 0.34 51.98± 2.88 0.00± 0.00 0.00± 0.00 0.00± 0.00Mst Kruskal 49.80± 1.54 61.66± 1.56 51.69± 1.53 0.00± 0.00 0.00± 0.00 0.00± 0.00Mst Prim 87.02± 0.90 95.33± 0.42 85.72± 1.56 6.67± 1.87 53.20± 3.03 4.47± 2.15Strongly Connected Components 94.45± 0.39 94.18± 0.04 96.19± 0.09 65.57± 0.75 64.40± 0.96 75.63± 1.76Topological Sort 94.30± 0.27 93.76± 0.01 98.56± 0.09 54.12± 0.78 52.02± 0.25 80.97± 0.67Mean 84.82 87.85 86.92 45.17 50.75 57.77metric (i.e., node-level metrics), we include a graph-level metric that is stricter thanthe node-level one. In particular, a graph is considered as correct if all of its nodesare classified correctly (in contrast to the node-level metric which measures the ac-curacy of the node classification). This way, we can better distinguish performancesamong models.We observe that while some algorithms struggle more on size generalization(e.g., DFS), others struggle on degree generalization (e.g., MST algorithms, FloydWarshall). This suggests that, even if we control all the distribution shifts causedby size generalization, distribution shift in degree is still challenging for somealgorithms and vice versa. Moreover, the graph-level metric shows very poorperformance in some tasks, despite a high node-level performance (e.g., FloydWarshall, and MST algorithms). This shows that in tasks with low graph-levelaccuracies, all graphs contain challenging parts. Finally, lower graph-level andnode-level accuracies of CLRS-Len-Deg compared to the other two datasets confirmthat when two distribution shifts are combined, the prediction problem becomesmore challenging.2.4 Improving OOD Generalization of NeuralAlgorithmic ReasoningIn this section, we introduce our solutions to some of the above challenges thatimprove the performance of end-to-end (i.e., without hints) neural algorithmic rea-soning. Since hints are disabled in our case, we make two changes to our model.21First, algorithms with the same input/output (e.g., Dijkstra and Bellman-Ford) onlydiffer in the number of processing steps. Since we discard these intermediatesteps, the number of steps becomes meaningless and we use a number of stepsindependent of the underlying algorithm. Moreover, the number of tasks shrinksfrom 30 to 24 tasks (e.g., Dijkstra and Bellman-Ford become Shortest Path, Quick-/Insertion/Merge/Bubble/Heap Sort become Sort, etc.). In order to facilitate thecomparison, we do not change the names of tasks and keep their original namesfrom the CLRS benchmark. Second, we use MPNN-G as the message-passingprocessor.2.4.1 Increase Dataset Size and RepresentativenessFirst, we propose to increase the dataset size. The purpose of this change is to trulymeasure the extrapolation properties of each model and avoid some in-distributiontraining issues like overfitting. We change the CLRS dataset configuration andincrease the amount of ID training samples from 1000 to 100K (and keep the IDvalidation and OOD test sets as they are). This way, it is possible to observe a muchmore diverse set of ID graphs. We note this strategy as a remedy to the low-dataregime and not a solution. This strategy especially becomes problematic and evenmore data-hungry as the size of training graphs grows, and sampling diverse graphsbecomes harder since the graph statistics become even more concentrated (e.g.,as the size of the graph grows, the number of samples to generate a disconnectedgraph increases). While this strategy is not a strong solution to OOD generalization,it helps mitigate any side factors contributing to low benchmark numbers but notdirectly related to size generalization; hence, providing a better test bed for testingsize generalization. Table 2.3 summarizes the different datasets we use in thischapter and their comparisons.2.4.2 Improving Representation Distribution ShiftThe position index of each node is a vital part of the inputs, especially since severalalgorithms depend on the ordering of such positions. Here, we explore a fewpossibilities of generalization to larger sizes. A suitable representation must benaturally extendable to any size of a graph. For instance, a one-hot encoding22Table 2.3: Summary of dataset configurations used in this paper.DatasetLengths Dataset SizeDataset GenerationTrain/Val Test Train Val/TestCLRS[94] 16 64 1000 32 Erdo˝s–Rényi graphs with fixed pL-CLRS 16 64 100000 32 Erdo˝s–Rényi graphs with fixed pL-CLRS-Len 16 32 100000 1000 K-Regular graphs with K = 4L-CLRS-Deg 32 32 100000 1000 K-Regular graphs with K = 4 for train/val, and K = 8 for testL-CLRS-Len-Deg 16 32 100000 1000 K-Regular graphs with K = 4 for train/val, and K = 8 for testof indices is not generalizable as it contains a fixed input length, and the inputdimension is not fixed across different sizes.As baselines, we consider the default index encoding used in the CLRS bench-mark (i.e., a scalar feature for each node i as (i− 1)/n where n is the number ofnodes). We also add positional encoding used by Transformers [89] as a baseline.We propose the following strategy as a mitigation to the distribution shift of thescalar index:Random Scalar Index. Although the Scalar Index representation is naturallyextendable to any graph size and node orderings are preserved, the model mightrely on spurious correlations between such numbers (e.g., the distribution shift dueto moving to larger sizes makes generalization harder). For instance, the modelmight rely on the fact that consecutive indices always have a difference of 1/nwhere n is the training dataset size, hence struggling to generalize as the size ofgraph grows and the difference between two consecutive indices becomes smaller.To mitigate the distribution shift, we propose to randomly sample scalars from[0,1] and use its sorted sequence during training, and use the same Scalar Indexfor validation and testing as mentioned before. This way, the model observes avariety of position embeddings and learns to better preserve the ordering of positions.In our experiments, we find that using random scalar makes training harder (e.g.,convergence is slower). As a remedy, we use a mix of both random and deterministicscalar index in a mini-batch as a training strategy to ease the training and achievea faster convergence (i.e., in each mini-batch, with the probability of 0.5 for eachgraph, we replace the position index of the whole graph with a random scalar index).In Appendix A.4, we also propose a more invariant (to shift and scale) way ofrepresentation, however, due to the superiority of the Random Scalar Index, we23Table 2.4: Performance of different positional encoding strategies on positions-order aware algorithms using MPNN-G processor on L-CLRS dataset.Method Scalar Random Scalar Position Enc.Dfs 25.36± 0.17 31.71± 1.13 11.90± 0.89Find Maximum Subarray Kadane 26.60± 5.95 47.85± 5.32 12.97± 3.96Dag Shortest Paths 99.82± 0.06 99.79± 0.02 31.45± 2.69Naive String Matcher 8.24± 2.33 7.67± 3.94 1.32± 0.35Matrix Chain Order 80.69± 1.33 92.07± 0.11 30.96± 3.16Topological Sort 82.62± 0.63 84.17± 1.67 33.28± 0.84Bfs 100.00± 0.00 100.00± 0.00 49.40± 4.83Mean 60.47 66.18 24.47choose Random Scalar for all remaining experiments.2.4.3 An Attention-Based Processor Competitive with MPNN-GInspired by recent works on higher-order message passing beyond node-level mes-sage passing, we explore this approach as a processor [14, 52, 66, 67]. We treateach edge of the input graph as a node, construct a new graph and run a GNN on topof that. Namely, the set of nodes, edges, and features V ′,E ′,X ′ will be as follows:V ′ = {v′i = (vi1 ,vi2) : vi1 ,vi2 ∈ V ,1≤ i≤ m}E ′ = {e′i = (ei1 ,ei2) : ei1 ,ei2 share a common node ,1≤ i≤ m′}X ′ = {x′i = [ x⊺vi1 , x⊺vi2 , x⊺ei ]⊺ : vi1 ,vi2 ∈ V ,ei ∈ E ,1≤ i≤ m}.In other words, nodes of the new graph are the edges of the original graph, edgesare between every two original edges that share a node, and node features areconcatenations of each endpoint of each edge. We use a transformer with edgemasks of E ′ as the GNN. Such a GNN architecture has a similar computationmechanism as a 2-Weisfeiler-Lehman (2WL) test, and hence we call this processor“2WL”. We, however, note that while the processors are limited by the expressivityof K-WL test [66], our model as a whole does not have such a limitation since in oursetting node features are equipped with a unique identifier [32]. The result of using24this processor is present in Table 2.5. As the results suggest, this processor showscompetitive results with MPNN-G, and outperforms MPNN-G in some algorithms,while under-perming MPNN-G in others. For instance, the 2WL processor showsa 20% higher performance on the String Matching problem, while the MPNN-Gshows a 20% higher performance on the Articulation Points finding algorithm.2.4.4 Hybrid ProcessorAs we saw above, different processors exhibit different generalization behaviors.Here, we show a hybrid processor brings performance gains over individual ones.To this end, we take MPNN-G and 2WL processors as they show competitive results,and are quite complementary in nature (e.g., max aggregation vs. attention, edge-wise message passing vs. node-wise message-passing). We consider the followingprocessor combination strategies:• Average: At each individual step of processing, two processors process thehidden representation, and the representation of each node/edge is computedusing an average of two processors. Namely,h(P)vi =h(P1)vi +h(P2)vi2and h(P)ei =h(P1)ei +h(P2)ei2,where P1, P2, P are two processors and the hybrid hidden representations.• Sigmoid: Similar to average, but a sigmoid is applied to the previous hiddendimension to assign a weighted combination of processor hidden representa-tion values. That is,h(P)vi = σ(h(P)vi W)h(P1)vi +(1−σ(h(En)vi W))h(P2)vih(P)ei = σ(h(P)ei W)h(P1)ei +(1−σ(h(En)ei W))h(P2)ei ,where W is a d×1 weight matrix.This processor is trained from scratch. Table 2.5 shows the results of experimentswith these processors. As the results suggest, a hybrid processor does bring an25Table 2.5: Mean test scores of different processors averaged over all 24 distinctalgorithms on the L-CLRS dataset (performances of individual algorithmsare averaged over three random seeds). The two rows correspond to usingthe scalar index and the random scalar index for input representationrespectively. The hybrid models are a combination of MPNN-G and 2WLwith the same parameter budget as individual processors (Table A.7 andA.8 contain the complete test scores of individual algorithms).Method MPNN-G 2WL Hybrid-Average Hybrid-SigmoidScalar Index 69.23 67.96 70.92 71.98Random Scalar Index 71.11 71.52 74.33 74.31Table 2.6: Summary of sequentially applying the improvements proposed inthis work. Test scores are averaged over all 24 distinct algorithms.Method Mean Test Score Improvement (%)Baseline 60.93 -Increase ID Dataset Size 69.23 8.30Improve Input Representation Shift 71.11 1.88Add Hybrid Processor 74.33 3.22average accuracy gain on performance. Furthermore, to show that the boost in theresults indeed originates from the complementariness of processors, we also reportthe results of a hybrid model of two processors of the same type (e.g., hybrid oftwo separate 2WL processors) in Table A.6 in the Appendix. Next, we combinethe proposed random scalar index representation and the 2WL processor in Section2.4.3 and report the average score on the OOD testing dataset in the bottom rowof Table 2.5. The hybrid models already outperform individual processors, andby adding the proposed input representation change, all processors gain 2%−3%improvement on the average accuracy. The hybrid processors achieve an averagescore of above 74% on the benchmark. A unified summary table of improvementsproposed in our work is shown in Table 2.6.262.5 Discussion on OOD PerformancesSo far, we have made a series of improvements to the models and training data andimproved the OOD accuracy of several tasks. Indeed, we have even achieved OODaccuracy above 98% on some tasks. Here, we ask “Is a task solved given our modelachieves a high score on the benchmark for that particular task?”. As discussed inSection 2.3.3, the test dataset might be biased and all graphs in the test set mighthave a common property and the model might take advantage of that property tooutput solutions. Following this issue, we pick the Bridges task on which the hybridof MPNN-G and 2WL achieves a high OOD performance as shown in Table A.8.As we show in Figure 2.4a, the bridges are all chain subgraphs branching out of aconnected component.To better evaluate models in this task, we create pairs of two-communitygraphs. Each pair of the community graph has two Erdo˝s–Rényi generated and inter-connected subgraphs. Then, we connect the two communities using a randomlychosen edge. We also connect the two communities of the second graph in the pairusing another random edge (see Figure 2.4b). This way, the first edge is a bridge inone graph and not a bridge in the second graph due to the existence of another edgeconnecting the communities. Ideally, a model capable of reasoning on bridges taskmust be able to identify the existence of the mentioned bridge in one graph, and theabsence of the bridge in the second graph of the pair. Although there might be morebridges among the pairs of graphs (i.e., inside individual communities), we do notconsider them in our metric since most of them might be trivial. Therefore, a randomguess gives 25% accuracy in such an evaluation. We evaluate the model on 1000randomly generated pairs of graphs, we find out that in only seven pairs of graphs,the model detects the existence of the mentioned bridge in both graphs correctly.This 0.70% accuracy is much lower than the 25% chance accuracy. This suggeststhat despite the remarkable improvement on the benchmark (originally from 72%to 98%), there is still a lot of room for improvement, and OOD generalization inthe Bridges task is not even close to being solved. Therefore, as the models getstronger, test datasets must also be refined to not only better evaluate the reasoningcapabilities of a model, but also prevent heuristics from achieving high scores.27(a) Erdo˝s–Rényi dataset (b) Two-Community datasetFigure 2.4: Illustration of two datasets generated for measuring the perfor-mance of a model on Bridges task. (a) The dataset generated by theCLRS benchmark. (b) We generate a two-community dataset for astricter measurement of the Bridges task. The dashed edge is present inone graph and absent in the other graph of the pair. Red edges show thetarget bridges. Green ovals show connected components. A model with98% OOD F1 score on the dataset (a) fails to detect the bridge betweenthe two communities of the dataset (b).2.6 Related WorksIn addition to the related works discussed in Section 2.3.1, there are some lines ofresearch close to this area that we explain below.Neural Algorithm Execution. A recent line of work has focused on the executionof algorithmic tasks using neural networks [1, 25, 49]. Some works considerNP-Hard problems over graphs and propose problem-specific solutions [16, 51],however, perfect extrapolation is infeasible for such problems [108]. Some worksconsider executing algorithmic tasks, however, they are either limited to simplearithmetic tasks [50] or are general purpose but hard-to-train and struggle to even fitin-distribution data for complex tasks [38]. [107] proposed a general neural executorto mimic elementary subroutines of any algorithm, but requires low-level designand assembly of subroutine engines for each specific task. More recent works haveproposed general-purpose GNN pipelines for executing (graph) algorithms [92, 93].Several benchmarks have been proposed to measure the reasoning capabilitiesof neural networks [94, 115, 117]. Perhaps the CLRS benchmark [94] is themost comprehensive benchmark to date, as it contains 30 highly complex andchallenging algorithmic reasoning tasks. In a concurrent work, [44] proposed aseries of improvements for improving OOD generalization when hints are used.28Expressivity of GNNs. Previous works have studied the expressiveness of GNNsfrom a WL-test point of view [67, 68, 100, 106]. These works show negative resultsby proving message passing GNNs are unable to distinguish two graphs if the WLtest fails to distinguish them. There exists a considerable body of literature forlifting such a limitation and improving the expressivity of GNNs by proposingnovel architectures and providing analysis [11, 13, 14, 63, 73]. Some prior workshave also addressed the limitations of GNNs by investigating their power in solvingalgorithmic tasks such as substructure counting in graphs [15, 22]. All theselimitations on the expressivity of GNNs, however, can be resolved by augmentingnode features [23, 32], at the cost of losing guaranteed equivariance. Since all tasksin our setting contain unique identifiers, none of the above limitations are applicableto our setting. This can also be verified by looking at validation scores. Here, thechallenge is to extrapolate to larger graphs, rather than the possibility of learning tosolve an in-distribution task.Reasoning Capabilities of Language Models. With the success of Large LanguageModels (LLMs) in natural language understanding tasks, a series of recent studieshave focused on extending the reasoning capability of LLMs even further. Thecore of these models is a Transformer architecture, which makes such tasks innature very close to neural algorithmic reasoning tasks. Recent works improve theperformance of LLMs using fine-tuning, chain-of-thought prompting [4, 55, 99],sampling multiple solutions [57, 97], and scaling [98]. The most relevant idea toour setting is chain-of-thought prompting, which forces the model to do multi-stepreasoning and is close in nature to the purpose of hints, but, in contrast to LLMs,hints do not bring any noticeable accuracy gain in neural algorithmic reasoning.Scaling the models might also be relevant, however, in our setting, the depth ofthe transformer must also increase with the size of the diameter of training graphs,which makes training large models on large graphs highly challenging.2.7 ConclusionIn this work, we investigate the OOD generalization of neural algorithmic reasoningtasks from an end-to-end training perspective. We identify and analyze severalchallenges of OOD generalization under the current framework. Then, we propose29solutions to some of these challenges by making the dataset larger, designingfixes to input representation distribution shift, and proposing an attention-basedprocessor. These improvements, lead to a final 74.33% accuracy averaged overall algorithms. Finally, we shed light on the limitation of our model in executingone of the benchmark tasks and conclude that despite our model gains a highOOD test accuracy, there still exists much room for improvement. Our analysisopens up several interesting directions for future work. Designing better hints withguaranteed invariance properties can potentially facilitate the step-by-step executionof algorithms and help improve OOD generalization. Sampling training graphsfrom a diverse set of distributions could help boost the generalization. Improvingthe model selection strategy and finding an alternative to the validation set couldhelp pick the best-performing models. Finally, it would be interesting to develop ameta-learning framework for inductive learning across graph sizes, in order for themodel to learn the invariance properties of an algorithm across different sizes andgeneralize to larger graphs.30Chapter 3Memorization Capacity ofMulti-Head Attention inTransformers3.1 IntroductionTransformers and Attention networks have achieved state-of-the-art results acrossdifferent domains such as natural language processing and vision tasks in recentyears [24, 85]. These models are typically very large, with billions of parameters,in order to achieve high performance. Given their size and complexity, an importantquestion is how well they can memorize the data they have been trained on. Con-cretely, the memorization capacity of a model fW corresponds to the minimal set ofparameters W such that fW (x) = y for all input-output pairs (x,y) in the dataset.Understanding the memorization capacity of a model is crucial to consider bothfrom a privacy perspective [20], and as a stepping stone towards quantifying themodel’s ability to generalize to new data [114].The question of memorization has been extensively analyzed for fully connectednetworks and dates back to the 1980s, when Baum [8] proved the first set of resultsfor single-hidden-layer Threshold networks with binary labels. The best knownbound on the memorization capacity of single-hidden-layer ReLU fully-connected31networks has only recently been established by Bubeck et al. [18] following nu-merous recent studies [110, 114]. Recent studies have also extensively consideredextensions to deeper models [72, 87, 95, 110]. However, the best known resultsconsist of order-wise bounds involving unkown constants, and logarithmic factors.Many of the above results on memorization of MLPs are established for data xi thatare assumed to be in General Position on continuous inputs. This is often a mildassumption in these settings since it is guaranteed for example under small randomperturbation.A Transformer consists of two main components: Multi-head Attention (MHA)module, and ReLU fully-connected network. Specifically, the MHA is the keycompent that differentiates the Transformer model from fully-connected networksas its operation involves Softmax combinations of its inputs that are calculated basedon inner products of the input representations. In practice, MHA successfully allowsthe Transformer model to jointly attend to information from different representationsand for processing long sequences. However, despite its prevalence in today’smodels, the theoretical understanding of the role of MHA is underexplored. Towardsclosing this gap, this chaptertakes a first step in studying the memorization capacityof MHA.Since compared to MLPs, MHA is (i) largely underexplored and (ii) fundamen-tally distinct in its underlyig mechanism, we take a step towards analyzing Attentiononly models by focusing on a single-layer MHA. The inputs to the MHA consist ofa n×d key matrix (also known as context) with n tokens, and a d×1 query vesctor.Each one of the H attention heads, computes a Softmax probablity vector based oninner product of linear emeddings of the key matrix and the query vectors, which isthen used to form a convex combination of the rows (aka tokens) of the key matrix.Then, the individual outputs of each head undergo a pooling operation to produce aprediction vector yˆ ∈ Rdout . These operations are parameterized in terms of severalweight matrices, of which, centeral role play in the kyey and query weight matricesW Kh,W Qh ∈ Rd×dh of each head h ∈ [H]. Given T inputs and corresponding labelsy(t) ∈ Rdout , t ∈ [T ], our goal is to show there exist weight matrices such that thepredictions yˆ(t) agree with the label vectors, i.e., yˆ(t) = y(t),∀t ∈ [T ].We explore how the number of attention heads of the model, and the contextsize of the input affect the memorization in this architecture. In particular, our32contributions are as follows:• We introduce two linear-independence based assumptions on the underlyingdataset. Our main introduced assumption is based on Kruskal Rank, whichis more relaxed than the General Poisition assumption, and encompassesGeneral Position as a specific case.• Under our introduced assumptions we prove that a single-layer MHA layerwith H heads, embedding dimension d, head dimension dh, context size n< d,output dimension dout ≤ d and O(Hd2) parameters, is capable of memorizingΘ(Hn). Under a configuration of dout = dh = d, and n = Θ(d), we showthe number of real-valued memorized numbers is optimal in the order ofparameters.• We verify our introduced assumption on real-world models. Specifically, weshow the Vision-Transformer (ViT) model on image classification satisfiesthe properties required by our assumptions partially in its embedding layer,and fully after its first layer. This is while the common assumption of GeneralPosition is violated at all levels.• By conducting experiments, we show that indeed memorization increaseslinearly with the number of heads, and monotonically with the context size,both of which confirm the results offered by our theorem.3.2 Related WorkTheoretical understanding of Transformers. Driven by the success of Transform-ers, there have been several attempts at understanding their underlying mechanisms,e.g., [28, 30, 46, 58, 60, 62, 80, 96]. A series of works have focused on the ex-pressivity of Transformers from a universal-approximation perspective [111, 112].However, these works typically show function approximation by constructing net-works with an exponential number of layers and parameters in input dimensions(context size and dimension ) to cover all the input space. In contrast, in memoriza-tion tasks, the aim is to find the most optimal network to achieve memorization ona finite number of samples. Bhojanapalli et al. [12], analyze multi-head Attention33expressivity on a single sequence of input and provide guidelines for setting thedimension of each head. However, to the best of our knowledge, our work is thefirst to explicitly address memorization in multi-head Attention with more than oneinput data point. Only very recently, Kim et al. [53] extend previous results forReLU networks [87] to Transformers with ReLU-activation feed-forward layer anda single Attention head and prove a bound with a sublinear number of parameters toachieve memorization. In our work, however, we isolate the effect of multi-headAttention for a better understanding of its role in the memorization capacity ofTransformers; thus, our results are not directly comparable as we do not incorporateany ReLU activation.Memory capacity of fully-connected networks. In contrast to Transformers,the memory capacity of fully-connected networks with ReLU/Threshold/Sigmoidnetworks has been extensively analyzed. Initially, Baum [8] provided the firstmemory bound for single hidden-layer Threshold networks. Under the GeneralPosition assumption, the author proved that a single hidden-layer Threshold canmemorize N binary labels with O(N) parameters. Bubeck et al. [18] extendedthe results to ReLU activation and real-valued outputs. Under such assumptions,these bounds are tight since memorization for a dense class of input data (e.g.,General Position) with piece-wise analytic definable functions (e.g., ReLU/Thresh-old/Sigmoid MLP) requires the number of parameters to be at least linear in theorder of the number of examples [82]. Subsequent works extended the results byconsidering different assumptions on the input data and considering multi-layernetworks [77, 95, 110]. More recent works on ReLU networks have providedbounds with O˜(√N) number of parameters w.r.t. number of data points by utilizingΩ˜(√N)-deep networks [72, 87]. Our findings and underlying assumptions on MHAclosely align with the principles established by Baum [8] and Bubeck et al. [18] inMLPs. We share a similar emphasis on the linearity of parameter growth and thenature of assumptions in our work.3.3 Problem SetupNotation. Throughout this chapter, we use the following notation: [n] denotesthe set of positive integers {1,2, . . . ,n}. Boldface upper case letters such as A34denote matrices. Boldface lowercase letters such as v denote vectors (columnvectors by default). We denote row-wise concatenation as [A;B], and column-wise concatenation as[A,B]. The set of n indexed vectors is denoted as x[n] ={x1,x2, . . . ,xn}.In this chapter, we consider a single-layer multi-head Attention (MHA) mecha-nism. An MHA layerA with H heads consists of three sets of matrices W Kh,W Qh ∈Rde×dh ,W V h ∈ Rde×dv for each head h ∈ [H], one weight matrix W O ∈ RHdv×de tocombine the outputs of different heads, and finally a read-out weight W D ∈ Rde×doutto get the final output of the model. We denote the entire set of parameters ofthis MHA layer as W = {{W Qh,W Kh,W V h}Hh=1,W O,W D}. For a key matrixE =[e⊤1 ;e⊤2 ; . . . ;e⊤n]∈ Rn×de containing n key vectors and a query vector e ∈ Rdeas input, this layer has the following computational mechanism:α h := EW KhW Q⊤h e (α h ∈ Rn) (3.1)θ h := Softmax(α h) (θ h ∈ Rn) (3.2)zh := E⊤θ h (zh ∈ Rde) (3.3)ph :=W V⊤h zh ∀ h ∈ [H] (ph ∈ Rdv) (3.4)o :=W O⊤[p⊤1 ; p⊤2 ; . . . ; p⊤H](o ∈ Rde) (3.5)yˆ :=W D⊤o (y ∈ Rdout), (3.6)where o is the output of the MHA layer and yˆ is the final output of the model. InTransformer terminology, the key matrix E is also called the context. Moreover,when e ∈ e[n], the MHA layer is called Self-Attention, and when the query vector isnot among key vectors, the MHA layer is a particular form of Cross-Attention. Inthis paper, we do not restrict the type of MHA as long as certain assumptions aresatisfied (see Section 3.4). Moreover, we define d := de.We consider a training set T consisting of T examples. Each example containsn context tokens, one query token, and one label for the query token. Namely, wehave T :={(E (t),e(t),y(t))}Tt=1. The goal is to find a parameter set W for theMHA layer such that the model is able to memorize any label set, i.e.,y(t) = fW (E (t),e(t)) for all t ∈ [T ].35For instance, consider image classification using Vision Transformer (ViT) [29].Two types of tokens are given as input to this model: A single “[CLS]” token, andseveral “image patch” tokens. The image patches contain information about theinput image. On the other hand, the “[CLS]” token is a special token and the classoutput of the model is read from the prediction of this token. For each image t ∈ [T ],the key matrix E (t) is the matrix containing image patches and the “[CLS]” token,the query vector e(t) corresponds to the “[CLS]” token only, and y(t) is the labelof the image. Note that, the type of MHA is here Self-Attention since the “[CLS]”token is inside E (t) as well, otherwise, the context would have contained only theimage patches.Often, a positional encoding is also added to the embedding of each token beforefeeding it to the Transformer module. Namely, for all t ∈ [T ],e(t) := x(t)+pos0, e(t)i := x(t)i +posi for all i ∈ [n],where x(t)i and x(t) are the raw embedding of input tokens, and posi ∈ Rd are fixedand unique positional encoding vectors for each position index. In the case of Self-Attention, the query vector takes only one positional encoding since it is alreadypresent among e[n].3.4 Main ResultsTo prove memorization bounds, two types of assumptions are typically made in priorworks: (1) norm-based assumptions [53, 72, 95], (2) linear independence assump-tions [8, 18]. In the former, the input data is assumed to be well-separated and/orhave a bounded norm, while in the latter, only some notion of linear independenceamong the data points is required. While norm-based assumptions are well-suitedfor discrete data such as language, linear independence is usually better satisfiedin continuous data such as images, or continuous distributions like Gaussian andUniform. For instance, Gaussianly distributed inputs satisfy linear independencewith probability one, while the separation and the bounded norm are violated withnon-zero probability.In this chapter, we focus on the second type of assumptions. Before stating theassumptions, we remind the reader of notions of Kruskal Rank and General Position.36Definition 1 (Kruskal Rank). Kruskal Rank of a set of vectors is the largest numberi such that every subset of size i is linearly independent.Definition 2 (General Position). A set of vectors {x1,x2, . . . ,xm} ⊂Rd is in GeneralPosition if it has the maximal Kruskal Rank of d.However, in contrast to previous works, we make a set of assumptions that differfrom General Position and are given as follows:Assumption 1. The set of all query vectors {e(t)|e(t) ∈Rd}Tt=1 has Kruskal Rank atleast n.Assumption 2. For each example t ∈ [T ], the context matrix E (t) ∈ Rn×d has rankn.In Section 3.5.1, we demonstrate concrete experimental settings for which theseassumptions hold while the General Position assumption fails. Before that, we firstmake use of these assumptions to state our main results below.Theorem 1. Assume that we have a multi-head attention layerA with H heads, em-bedding dimensions d := de, dv≥ dout≥ 1, and dh≥ 1. LetT ={(E (t),e(t),y(t))}Tt=1be a training set with context size n< d. Define n1 := min(n,dh). If the set of queryvectors {e(t)}Tt=1 has Kruskal Rank at least n, E (t) has rank n for each examplet ∈ [T ], andT ≤ H (n1−1)+1,then, there exists a set of parameters W such that A that can memorize T .Theorem 1 shows that an MHA with 2Hddh +Hddv + ddout parameters canmemorize at least H (min(n,dh)−1)+1 examples under Assumptions 2 and 1. Togain further insights on the implications of this theorem consider three concretesettings: (1) Suppose dh = d and scalar model output dv = dout = 1, then a MHAwithΘ(Hd2) can memorizeΩ(Hn) examples. (2) Suppose dh = n< d,dv = dout = 1then a MHA withΘ(Hdn) can memorizeΩ(Hn) examples. (3) Suppose dh = dv = dand d-dimensional model output dout = d, then a MHA with Θ(Hd2) can memorize37Θ(Hn) examples with d-dimensional label each. The third scenario is particularlyinteresting. To see this further, consider the typical setting where n and d are ofthe same order (n = Θ(d)), then the established bound becomes optimal. Thisoptimality is reflected in the fact that Ω(Hnd) = Ω(Hd2) real-valued outputs arememorized, and a trivial upper bound for the memorization of a MHA is the numberof its parameters, which is Θ(Hd2).The assumption that n < d is general practice. For instance, in prominentmodels such as Vision Transformer [29], BERT [27], and GPT-3 [17], the n/dratio is 197/768, 512/768, and 2048/12288, respectively. Furthermore, in Sections3.4.3 and 3.5.1, we demonstrate that Assumption 2 typically holds in practice dueto positional encoding, and Assumption 1 is valid after applying a single Attentionlayer. Unlike previous works on MLPs that regard General Position as an assumption[8, 18], our approach innovatively incorporates Assumption 1, which imposes aweaker constraint than General Position since n< d. In fact, our experiments revealthat Assumption 1 is satisfied in real data, even though the more commonly usedGeneral Position assumption is often not met.3.4.1 Proof of Theorem 1Our proof consists of two steps: (1) We focus on the intermediate representationsof the network that originate from zh in Equation (3.3). We demonstrate thatwith a suitable set of attention weights {(W Kh,W Qh)}Hh=1, we can achieve linearindependence across different data points, which results in a rich representation. (2)Then, we prove the existence of {{W V h}Hh=1,W O,W D} for achieving memorizationby taking advantage of the rich representation from the previous step.Let us define matrix Z as follows:z(t) :=[z(t)⊤1 ,z(t)⊤2 , . . . ,z(t)⊤H]⊤ ∈ RdH (3.7)Z :=[z(1)⊤;z(2)⊤; . . . ;z(T )⊤]∈ RT×dH . (3.8)Note that Z is a function of input data, and key/query weight matrices W Kh,W Qh.Each vector z(t) is a concatenation of H d-dimensional vectors each representing aconvex combination of context vectors of the t-th example parameterized in terms38of key/query vectors and W Kh,W Qh. The matrix Z gathers all these dH dimensionalrepresentations, one for each example in a matrix form. Clearly, this matrix is afunction of the input data and of the key/query matrices W Kh,W Qh,h ∈ [H]. Then,based on the computation mechanism of the MHA layer, the predicted labels yˆ(t)are given byyˆ(t) := Z︸︷︷︸Step 1W V 1W V 2. . .W V HW OW D︸ ︷︷ ︸Step 2∈ R, ∀t ∈ [T ]. (3.9)In the first step, using an inductive proof technique, we show that Z has a high rank,and use this information in the second step to solve a linear system of equations andachieve memorization.Step 1. The Rank of Z. Our key technical result is proving the following lowerbound on the rank of Z .Lemma 1. Under the conditions of Theorem 1, there exists parameters {(W Kh,W Qh)}Hh=1such thatrank(Z)≥min{H(n1−1)+1,T} .Lemma 1 proves the existence of weights W Kh,W Qh for achieving linear inde-pendence on the rows of Z . We provide a proof sketch of this lemma here and deferthe full proof to Appendix B.1. The main idea behind proving the lower bound is toinductively add heads and use each new head for a distinct set of n1−1 examples,which in turn translates to an increase on the ranl of Z . The new head exploits a newdistinct set of examples without affecting the rank of the previous induction step.As we will see, to accomplish this, our idea is to leverage the saturation property ofthe softmax nonlinearity.In order to get H(n1− 1)+ 1 rank, each time a head is added, the rank of Zshould be increased by n1− 1. For each head h, we only focus on h(n1− 1)+ 1examples, therefore, the subsequent step can be viewed as adding a new head withtunable weights, and a set of n−1 fresh data points. Namely, the induction step of39Z can be seen as the following block matrix:Z =Z ′a Z′′aZ ′b Z′′b︸ ︷︷ ︸hd︸︷︷︸d .Above, Z ′a ∈ R(T−n+1)×hd is the matrix derived from the induction hypothesis,Z ′b ∈ R(n−1)×hd contains the embeddings of new examples corresponding to theheads of the previous induction steps, Z ′′a ∈ R(T−n+1)×d contains the embeddings ofthe previous examples corresponding to the new head, and finally, Z ′′b ∈ R(n−1)×dcontains the embeddings of the new examples corresponding to the new heads. Wetune the new head’s weights in a way that Z ′′b would increase the rank of Z . However,by such a tuning, the whole right block changes, and increasing the rank becomeschallenging. To overcome this challenge, we exploit the saturation property ofSoftmax to keep Z ′′a fixed while tuning Z′′b and increasing the rank of Z .Step 2. Solving system of equations. Next, we prove we can find the remain-ing weights to memorize at least rank(Z) examples. WLOG, let us assume thatT = rank(Z), otherwise we can ignore some of the data points and keep onlyrank(Z) data points that make Z full-rank. We find weights {{W V h}Hh=1,W O,W D}to memorize all T examples as follows. Since Z has full row rank, there exist weightmatrix W ∈ RHd×dv such thatZW =y(1)⊤y(2)⊤... 0T×(dv−dout)y(T−1)⊤y(T )⊤.40has a solution. Now, we can set the remaining weights as followsW V 1W V 2...W V H=W , W O =Idv×dIdv×d...Idv×d , W D = Id×dout . (3.10)Substituting in Eq. (3.9), this gives yˆ(t) = y(t) for all t ∈ [T ]. Therefore, we haveproved that the memorization capacity of one layer is at least rank(Z). But, sincealso we proved a lower bound of H(n1−1)+1 on the rank of Z , we conclude thatthe memorization capacity is at least H(n1−1)+1, as advertised.Remark 1. In the setting of Theorem 1, if Assumption 1 is violated (i.e., the KruskalRank of query vectors is less than n), a weaker lower bound on the memorizationstill holds:T ≤ H(min(Q,dh)−1)+1,where Q is the Kruskal Rank of query vectors {e(t)|e(t) ∈ Rd}Tt=1.3.4.2 How tight is the lower bound on the rank?As seen above, our proof of Theorem 1 relies on the lower bound of rank(Z)in Lemma 1. Hence, one might naturally wonder how tight this bound is. Theproposition below proves that our bound of Lemma 1 is indeed tight.Proposition 1. Under the conditions of Theorem 1, if all the contexts are shared,i.e., E (t) := E for all t ∈ [T ], then for any {(W Kh,W Qh)}Hh=1, we haverank(Z)≤ H(n−1)+1.Proof. First, note that when the contexts are shared, the matrix Z has the followingdecomposition:Z⊤ =(IH ⊗E T)Θ⊤,where ⊗ is the Kronecker product and Θ ∈ RT×Hn is defined similar to Z (i.e., theblock in t-th row and h-th column is θ (t)⊤h ). Notice that Θ ∈ RT×Hn, therefore, the41rank of Z can be at most Hn. However, we prove here that the column rank of Θ isat most H(n−1)+1 due to the structure of θ (t)h being convex combination vectors.With a slight abuse of notation, let us write Θ as:Θ =[Θ1, Θ2, Θ3 . . . , ΘH],where each Θh ∈ RT×n corresponds block belonging to head h ∈ [H]. Observe thatΘh is a right stochastic matrix (i.e., rows sum up to one) for all h ∈ [H]. Thus, eachrow of Θh−Θ1 sums up to zero and the matrix has at least one zero eigenvalue.Therefore,rank(Z)≤ rank(Θ)= rank([Θ1, Θ2, Θ3, . . . , ΘH])= rank([Θ1, Θ2−Θ1, Θ3−Θ1, . . . , ΘH −Θ1])≤ rank(Θ1)+ rank(Θ2−Θ1)+ · · ·+ rank(ΘH −Θ1)≤ n+(n−1)+(n−1)+ · · ·+(n−1)= H(n−1)+1 .Hence rank(Z) is upper bounded by H(n−1)+1 and this concludes the proof.Proposition 1 shows that the lower bound Lemma 1 on the rank of Z is tightwhen the context is shared across data points and dh ≥ n.Furthermore, the setting of Proposition 1 provides an elegant worst-case scenariofor Assumptions 1 and 2. For instance, consider this setting for n = 1, i.e., eachcontext contains a single token and is shared across all examples. Then, all θ (i)hbecome scalars, whose values are equal to one. Consequently, the vector z(t)h copiesthe first (and only) context token, which is shared across all examples. As a result,no matter how large H or d are, the output of MHA for all examples is identical,and the memory capacity of the attention is exactly T = H(min(n,dh)−1)+1 = 1(i.e., a single example), which agrees with both Theorem 1 and Proposition 1.423.4.3 Analysis of the validity of assumptionsWe verify that our assumptions either hold on the embedding layer or hold byapplying a single Attention layer. First, we examine Assumption 2. This assumptionstates that inside each example, the context matrix is full-rank. In most tasks, apositional encoding is added to each token representation before feeding the datainto the Transformer architecture. One of the widely used positional encodingschemes is sinusoidal positional encodings [89], where the positional encodingmatrix is full rank, which in practice, makes context matrix E (t) full-rank as well.We also verify this full-rank assumption in Section 3.5.1 on real data.Next, we examine Assumption 1. Let us consider the case of image classificationusing ViT. Recall that the query vectors e(t), in this case, are all equal to the “[CLS]”token plus positional encoding of position zero. Therefore, they are all equal, andAssumption 1 does not hold. However, the context for each image is different.Intuitively, the role of “[CLS]” token is to mix the information from the context. Infact, we prove under mild assumptions, even a single-head Self-Attention Layerwith a skip connection can achieve token mixing. Hence, when looking at the outputof the (even randomly initialized and fixed) Self-Attention layer, the assumptions ofTheorem 1 hold, and our theorem can be applied to the second layer of Attention toachieve memorization. This argument is similar to arguments by Vershynin [95],who first designs layers with random weights to enrich the representation of datapoints, and then proves the existence of a learned subsequent perception layer formemorizing the labels. We elaborate on this below.First, we formally explain how a Self-Attention layer with skip connectionworks. Consider a single example with context matrix E and query token e. Given aset of attention weightsW , the Self-Attention layer with skip connection transformsthe tokens according to the following equations:e′ := e+ fW(E ,e)e′i := ei+ fW(E ,ei)for all i ∈ [n]E ′ :=[e′⊤1 ;e′⊤2 ; . . . ;e′⊤n],where e′, and E ′ correspond to the output of query, and context vectors, respectively.43Note that the query matrix is part of the context matrix due to Self-Attention. Letus assume the query vector is in the first row of E , i.e. e1 = e, and consequentlye′ = e′1.Remark 2. The results of Theorem 1 also hold true for Attention layers with skipconnection (see Appendix B.3).We are now ready to state a proposition on the mixing power of this type ofAttention with skip connection. First, let us introduce a new relaxed assumption.Assumption 3. For a training dataset T ={(E (t),e(t),y(t))}Tt=1, the set S de-fined ase˜(t) := e(t)+1n(e(t)1 + e(t)2 + · · ·+ e(t)n)for all t ∈ [T ], (3.11)S :={e˜(1), e˜(2), . . . , e˜(T )}, (3.12)has Kruskal Rank at least n.Proposition 2. Define a Self-Attention layer A0 with a single head, dv = d param-eters, weight matrices W K =W Q = 0,W V =W O = I , and a skip connection. LetT ={(E (t),e(t),y(t))}Tt=1be a training set under Assumptions 2 and 3. Moreover,let E ′(t),e′(t) be the output of A0 for each example t ∈ [T ] and define a new datasetT ′ asT ′ :={(E ′(t),e′(t),y(t))}Tt=1.Then, both Assumptions 2 and 1 hold true for T ′.Assumption 3 assumes that the input data as a whole is sufficiently distinctiveacross different examples. For instance, in the case of images, it corresponds to theaverage of all image patches (plus a constant vector) being sufficiently distinctiveacross different images. Proposition 2 then proves under such assumptions (i.e.,Assumption 3 together with Assumption 2), by applying a trivial Self-Attentionlayer, we get Assumption 1 at the output of the Attention layer while maintainingAssumption 2. Therefore, the “[CLS]” token can perform token mixing even with atrivial Attention. This means that Theorem 1 can be applied to a second layer ofAttention to memorize the examples.443.5 Experiments3.5.1 Testing assumptions on Vision Transformer (ViT)As mentioned in Section 3.4.3, we show that our assumptions on the input tokenshold after applying a single Attention layer. To test this, we mainly focus on imageclassification and Vision Transformer (ViT) [29, 102] to conduct our experiments.Namely, we consider the following models:• Embedding Layer: A randomly initialized embedding layer. This layercontains only positional encoding and a linear embedding layer (i.e., the inputtokens E (t),e(t)).• Random Attention: A randomly initialized Self-Attention layer with sim-ilar architecture to a ViT layer, without the Multi-Layer Perceptron of aTransformer.• Random ViT: A randomly initialized base ViT. We only look at the output ofthe first layer and discard the subsequent layers.• Trained ViT: Similar to the random ViT, but with the difference that we takethe weights from the ViT pre-trained on ImageNet [26].We evaluate the mentioned models on 2000 images sampled from ImageNet.To empirically test Assumption 2, we verify whether the context vectors are alllinearly independent for each example. On the other hand, testing Assumption 1 iscomputationally difficult since computing Kruskal Rank is NP-Hard. Therefore, itcan only be approximately tested in polynomial time. To address this, we randomlysample n query vectors from the training set T and construct a matrix of size n×d.We then test if the rank of the matrix is equal to n. After repeating this test 5000times, we report that the assumption holds if the rank check holds for at least 99%of the test instances. The results of our empirical tests are summarized in Table3.1. Specifically, we find that only Assumption 1 is violated in the Embeddinglayer, whereas both assumptions hold true for the remaining models. Notably, thecommonly made assumption of General Position fails in all cases, and the KruskalRank is typically much less than d, as shown by Figure 3.1. We also empirically test450 200 400 600 800Kruskal Rank Candidate0.00.20.40.60.81.0Satisfied FractionndFigure 3.1: Testing Kruskal Rank of query tokens on the output of one layerRandom Attention on ImageNet dataset. The Kruskal Rank is onlyslightly larger than n (Assumption 1), and much smaller than d (GeneralPosition).Table 3.1: Testing assumptions on ImageNet and ViT: while only one assump-tion is violated in the embedding layer, both hold true after a single layerof Attention or Transformer. The General Position assumption is violatedin all cases.Model Embedding Rand Attention Rand ViT Trained ViTGeneral Position × × × ×Assumption 1 × ✓ ✓ ✓Assumption 2 ✓ ✓ ✓ ✓Assumption 3 and find that it already holds for the Embedding layer. This may beexplaining why Assumption 1 starts to hold after applying any type of Attention.3.5.2 Memory capacity increases with the number of headsTheorem 1 provides two conclusions: (1) When fixing dimension d, increasingthe number of heads H improves memorization. (2) When fixing dimension d,and number of heads H, increasing the context size n improves memorization. Byconducting synthetic experiments, we verify the conclusions drawn by the theorem.In our experiments, we fix the dimension d = 64. To strictly follow Assumptions2 and 1, we generate random query inputs, and random shared context inputs (seeAppendix B.5.1 for the non-shared context experiments) according to a uniform46distribution, i.e.,x1,x2, . . . ,xniid∼U (0,1)64x(t) iid∼U (0,1)64 for all t ∈ [T ]x(t)i := xi for all i ∈ [n], t ∈ [T ],where x(t)i ,x(t) are the raw inputs. Then each input is given to the model whichconsists of an embedding layer, positional encoding, and an Attention layer. To testthe first conclusion, we fix n = 32 and increase H. To test the second conclusion,we fix H = 20 and increase context size n. In addition to regression tasks, we alsoinclude classification tasks since measuring memorization in classification tasks ismore straightforward (i.e., the accuracy metric is more direct). For regression tasks,we generate uniform labels y(t) ∼ U (0,1), and for classification tasks, we drawuniform labels from 100 classes y(t) ∈U ([100]). We measure the memorizationfor the classification task by taking the average accuracy across examples inside adataset. As for the regression tasks, we measure the mean squared error (MSE) lossto compare memorization. We train each task for at least 500 epochs with a batchsize of 256, using an Adam optimizer with a learning rate of 0.001, and a schedulerwith linear warmup and Cosine decay.The results of our experiments are presented in Figure 3.2 for both classificationand regression tasks and various dataset sizes T . As demonstrated by the figure,increasing H and n not only increases memorization performance but also it does soin a linear fashion, both of which confirm the results suggested by Theorem 1 andProposition 1.While our experimental results conclusively show that increasing the numberof heads H or context size n improves the model’s memorization capability, wecaution the reader that the experiments are done purely through gradient descentoptimization, which has limitations when it comes to memorization. First, it is notguaranteed that if a memorizing solution exists, GD would be able to find it. Second,optimization requires hyperparameter tuning and a large number of gradient updatesto achieve memorization, especially in under-parameterized regimes and large-context sizes [35]. We mitigate these issues by running small-scale experiments470 10 20 30 40 50 60# Heads (H)0.00.20.40.60.81.0Memorized FractionClassification AccuracyT=40kT=30kT=20kT=10k(a)0 10 20 30 40 50 60# Heads (H)10 1010 810 610 410 2MSE LossRegression LossT=30kT=25kT=20kT=15kT=10k(b)0 10 20 30 40 50 60Context Size (n)0.20.40.60.81.0Memorized FractionClassification AccuracyT=18kT=14kT=10kT=6k(c)0 10 20 30 40 50 60Context Size (n)10 910 710 510 310 1MSE LossRegression LossT=25kT=20kT=15kT=10k(d)Figure 3.2: Testing memorization as a function of heads (a, b) and context size(c, d) for both classification and regression, under Assumptions 2 and 1.Synthetic examples with d = 64 and shared context (see Proposition 1).Both settings show a positive correlation between heads/context size withmemorization capacity, while the classification setting shows a linearrelationship with both. The variance in regression plots is magnified insome cases due to the log scale of the plot.with a sufficiently high number of optimization steps and adequate learning ratescheduling. Running the experiments reported in our paper takes approximately10 GPU days on an NVIDIA V100 GPU. However, this procedure becomes highlyresource exhaustive and prohibitive in larger-scale experiments.483.6 Conclusion and DiscussionIn this chapter, we proved a lower bound on the memorization capacity of attentionlayers under a rather mild linear independence assumption. Our proof techniquerelies on establishing a tight bound on the rank of a matrix of intermediate represen-tations in the attention layer. Specifically, we showed that an attention layer withΘ(d2H) parameters can memorize Ω(nH) examples, where n is the context size andH is the number of attention heads. This bound holds as long as the assumptions oflinear independence between query vectors, and full rank for context matrices ofeach example hold. We demonstrated that these assumptions are typically satisfiedin practice by providing experiments on image classification using ViT. Finally,using synthetic experiments we showed a linear relationship between memorizationand the number of heads.Our work opens up several promising avenues for future research. First, it wouldbe valuable to extend our theoretical results beyond one layer of attention networks.Second, establishing upper bounds on the memorization capacity of MHA remainsan important open question. Finally, incorporating additional assumptions to narrowthe gap between practical memorization abilities and theoretical guarantees couldshed more light on the memorization capacity of Attention networks (see AppendixB.5). Ultimately, we hope that gaining a better understanding of the memorizationabilities of Transformers could facilitate more efficient implementations and bet-ter use of data for training more effective, generalizable, and privacy-preservingmodels.49Chapter 4Conclusion and Future Work4.1 ConclusionIn this thesis, we investigated the challenges of achieving out-of-distribution gener-alization and true reasoning abilities in neural networks, particularly Transformers.We analyzed the issues that arise in neural algorithmic reasoning and proposedsolutions to address some of these challenges. Our proposed improvements leadto a 13.4% improvement over the previous state-of-the-art results. Despite all theimprovements, and the models achieving high in-distribution accuracy, there is stilla significant gap between ID and OOD performance, which signals the lack ofreasoning.Motivated by this issue, we then took a step back and analyzed the expressivityof Transformers. As a stepping stone in this direction, we provided a theoreticalanalysis of the memorization capacity of a single-layer MHA. Overall, our analysisestablished a linear relationship of capacity with the number of attention heads,which we also confirm experimentally together with the validity of our linear-independence assumptions.4.2 Future WorkOur work opens up several interesting future directions. The following futuredirections are particularly interesting to pursue.504.2.1 Improving memorization bounds for TransformersIn Chapter 3 we proved memorization bounds for a single layer MHA and a singleoutput for each example in the dataset. Extending the analysis to multi-layer MHAand examining the role of depth in these models will provide valuable insights intothe role of depth in memorization capability. Furthermore, we plan to broaden ourscope beyond single outputs. Exploring sequence-to-sequence MHA models shedslight on the role of MHA inside each sequence. Finally, understanding the interplaybetween the MHA and MLP components of Transformers is of particular interest.This exploration will offer a deeper understanding of how these two fundamentalelements cooperate in the memorization process.4.2.2 Closing the gap between theoretical and empirical memorizationpower of TransformersIn Chapter 3 we showed a linear relationship between memorization capacityand number of parameters. However, in practice, neural networks show superiormemorization capabilities on real-world data. While we took an initial step towardsverifying our assumptions input data, it would be interesting to introduce furtherrealistic assumptions to better explain empirical memorization in Transformers.For instance, in real-world data, the labels are not chosen arbitrarily, and thereexist particular patterns both across examples and between each example and itslabel. Therefore, by introducing further assumptions to exploit such patterns, wemay not only establish connections between memorization and generalization butalso predictively quantify memorization in real-world applications such as naturallanguage processing.4.2.3 Establishing connection between expressivity and reasoningThe level of reasoning (e.g., arithmetic reasoning) in current language modelssignificantly improves with the scale of the model (i.e., larger number of parameters).This improvement in some cases manifests abruptly and emerges at a certain scalethreshold, and does not exist below a certain threshold [98]. It would be interestingto investigate how such higher expressivity contributes to improved reasoning inmodels and establish a connection between memorization and reasoning.514.2.4 Improving reasoning performance of models beyondarchitectural modificationsIn Chapter 2, we introduced a range of enhancements, including architectural ad-justments, aimed at augmenting performance on the CLRS algorithmic reasoningbenchmark. Although we extensively investigated diverse architectural modifica-tions, there remains untapped potential within other facets of the training process.To illustrate, a mere expansion in dataset size resulted in a remarkable improve-ment in reasoning accuracy. Exploring more refined sampling techniques to bettercover the underlying function holds promise for yielding even more substantialadvancements in reasoning capability. For instance, finding a minimal subset ofexamples as a training dataset to cover all edge cases of a reasoning algorithm couldfacilitate the learning algorithm to find more generalizable solutions. Furthermore,improving the training objective itself could potentially elevate accuracy whendealing with out-of-distribution (OOD) data. Specifically, drawing inspiration fromthe literature in curriculum learning [2] and meta-learning [33] might offer valuableinsights into objectives and training procedures to enhance both learning speed andthe reasoning abilities of models.52Bibliography[1] E. Abbe, S. Bengio, E. Cornacchia, J. M. Kleinberg, A. Lotfi, M. Raghu, andC. Zhang. Learning to reason with neural networks: Generalization, unseendata and boolean measures. ArXiv, abs/2205.13647, 2022. → page 28[2] E. Abbe, E. Cornacchia, and A. Lotfi. Provable advantage of curriculumlearning on parity targets with mixed inputs. ArXiv, abs/2306.16921, 2023.URL https://api.semanticscholar.org/CorpusID:259286830. → page 52[3] M. C. Angelini and F. Ricci-Tersenghi. Cracking nuts with a sledgehammer:when modern graph neural networks do worse than classical greedyalgorithms. ArXiv, abs/2206.13211, 2022. → page 17[4] C. Anil, Y. Wu, A. J. Andreassen, A. Lewkowycz, V. Misra, V. V. Ramasesh,A. Slone, G. Gur-Ari, E. Dyer, and B. Neyshabur. Exploring lengthgeneralization in large language models. ArXiv, abs/2207.04901, 2022. →pages 14, 29[5] M. Arjovsky, L. Bottou, I. Gulrajani, and D. Lopez-Paz. Invariant riskminimization. ArXiv, abs/1907.02893, 2019. → page 14[6] P. Bartlett, D. J. Foster, and M. Telgarsky. Spectrally-normalized marginbounds for neural networks, 2017. → page 7[7] P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzalez, V. Zambaldi,M. Malinowski, A. Tacchetti, D. Raposo, A. Santoro, R. Faulkner, et al.Relational inductive biases, deep learning, and graph networks. arXivpreprint arXiv:1806.01261, 2018. → page 11[8] E. B. Baum. On the capabilities of multilayer perceptrons. Journal ofComplexity, 4(3):193–215, 1988. ISSN 0885-064X.doi:https://doi.org/10.1016/0885-064X(88)90020-9. → pages7, 31, 34, 36, 3853[9] Y. Bengio, N. Léonard, and A. C. Courville. Estimating or propagatinggradients through stochastic neurons for conditional computation. ArXiv,abs/1308.3432, 2013. → page 71[10] B. Bevilacqua, Y. Zhou, and B. Ribeiro. Size-invariant graph representationsfor graph classification extrapolations. In ICML, 2021. → page 15[11] B. Bevilacqua, F. Frasca, D. Lim, B. Srinivasan, C. Cai, G. Balamurugan,M. M. Bronstein, and H. Maron. Equivariant subgraph aggregation networks.ArXiv, abs/2110.02910, 2022. → page 29[12] S. Bhojanapalli, C. Yun, A. S. Rawat, S. Reddi, and S. Kumar. Low-rankbottleneck in multi-head attention models. In H. D. III and A. Singh, editors,Proceedings of the 37th International Conference on Machine Learning,volume 119 of Proceedings of Machine Learning Research, pages 864–873.PMLR, 13–18 Jul 2020. → page 33[13] C. Bodnar, F. Frasca, N. Otter, Y. G. Wang, P. Lio’, G. Montúfar, and M. M.Bronstein. Weisfeiler and lehman go cellular: Cw networks. In NeurIPS,2021. → page 29[14] C. Bodnar, F. Frasca, Y. G. Wang, N. Otter, G. Montúfar, P. Lio’, and M. M.Bronstein. Weisfeiler and lehman go topological: Message passingsimplicial networks. ArXiv, abs/2103.03212, 2021. → pages 24, 29[15] G. Bouritsas, F. Frasca, S. Zafeiriou, and M. M. Bronstein. Improving graphneural network expressivity via subgraph isomorphism counting. IEEEtransactions on pattern analysis and machine intelligence, PP, 2022. →page 29[16] X. Bresson and T. Laurent. The transformer network for the travelingsalesman problem. ArXiv, abs/2103.03012, 2021. → page 28[17] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal,A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal,A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler,J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess,J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei.Language models are few-shot learners. In H. Larochelle, M. Ranzato,R. Hadsell, M. Balcan, and H. Lin, editors, Advances in Neural InformationProcessing Systems, volume 33, pages 1877–1901. Curran Associates, Inc.,2020. URL https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf. → pages 1, 6, 3854[18] S. Bubeck, R. Eldan, Y. T. Lee, and D. Mikulincer. Network size and size ofthe weights in memorization with two-layers neural networks. In NeuralInformation Processing Systems, 2020. URLhttps://api.semanticscholar.org/CorpusID:227275586. → pages7, 32, 34, 36, 38[19] S. Bubeck, V. Chandrasekaran, R. Eldan, J. Gehrke, E. Horvitz, E. Kamar,P. Lee, Y. T. Lee, Y. Li, S. Lundberg, H. Nori, H. Palangi, M. T. Ribeiro, andY. Zhang. Sparks of artificial general intelligence: Early experiments withgpt-4, 2023. → page 1[20] N. Carlini, F. Tramèr, E. Wallace, M. Jagielski, A. Herbert-Voss, K. Lee,A. Roberts, T. B. Brown, D. X. Song, Ú. Erlingsson, A. Oprea, and C. Raffel.Extracting training data from large language models. In USENIX SecuritySymposium, 2020. → pages 7, 31[21] Y. Chen, Y. Zhang, Y. Bian, H. Yang, K. Ma, B. Xie, T. Liu, B. Han, andJ. Cheng. Learning causally invariant representations for out-of-distributiongeneralization on graphs. In Advances in Neural Information ProcessingSystems, 2022. → page 14[22] Z. Chen, L. Chen, S. Villar, and J. Bruna. Can graph neural networks countsubstructures? ArXiv, abs/2002.04025, 2020. → page 29[23] G. Dasoulas, L. D. Santos, K. Scaman, and A. Virmaux. Coloring graphneural networks for node disambiguation. In IJCAI, 2020. → page 29[24] M. Dehghani, J. Djolonga, B. Mustafa, P. Padlewski, J. Heek, J. Gilmer, A. P.Steiner, M. Caron, R. Geirhos, I. Alabdulmohsin, R. Jenatton, L. Beyer,M. Tschannen, A. Arnab, X. Wang, C. Riquelme Ruiz, M. Minderer,J. Puigcerver, U. Evci, M. Kumar, S. V. Steenkiste, G. F. Elsayed,A. Mahendran, F. Yu, A. Oliver, F. Huot, J. Bastings, M. Collier, A. A.Gritsenko, V. Birodkar, C. N. Vasconcelos, Y. Tay, T. Mensink,A. Kolesnikov, F. Pavetic, D. Tran, T. Kipf, M. Lucic, X. Zhai, D. Keysers,J. J. Harmsen, and N. Houlsby. Scaling vision transformers to 22 billionparameters. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato,and J. Scarlett, editors, Proceedings of the 40th International Conference onMachine Learning, volume 202 of Proceedings of Machine LearningResearch, pages 7480–7512. PMLR, 23–29 Jul 2023. URLhttps://proceedings.mlr.press/v202/dehghani23a.html. → pages 1, 3, 3155[25] G. Del’etang, A. Ruoss, J. Grau-Moya, T. Genewein, L. K. Wenliang,E. Catt, M. Hutter, S. Legg, and P. A. Ortega. Neural networks and thechomsky hierarchy. ArXiv, abs/2207.02098, 2022. → page 28[26] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. Imagenet: Alarge-scale hierarchical image database. In 2009 IEEE conference oncomputer vision and pattern recognition, pages 248–255. Ieee, 2009. →page 45[27] J. Devlin, M. Chang, K. Lee, and K. Toutanova. BERT: pre-training of deepbidirectional transformers for language understanding. In J. Burstein,C. Doran, and T. Solorio, editors, Proceedings of the 2019 Conference of theNorth American Chapter of the Association for Computational Linguistics:Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA,June 2-7, 2019, Volume 1 (Long and Short Papers), pages 4171–4186.Association for Computational Linguistics, 2019. doi:10.18653/v1/n19-1423.URL https://doi.org/10.18653/v1/n19-1423. → pages 38, 90[28] Y. Dong, J.-B. Cordonnier, and A. Loukas. Attention is not all you need:pure attention loses rank doubly exponentially with depth. In M. Meila andT. Zhang, editors, Proceedings of the 38th International Conference onMachine Learning, volume 139 of Proceedings of Machine LearningResearch, pages 2793–2803. PMLR, 18–24 Jul 2021. URLhttps://proceedings.mlr.press/v139/dong21a.html. → page 33[29] A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai,T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly,J. Uszkoreit, and N. Houlsby. An image is worth 16x16 words: Transformersfor image recognition at scale. In 9th International Conference on LearningRepresentations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021.OpenReview.net, 2021. URL https://openreview.net/forum?id=YicbFdNTTy.→ pages 1, 3, 4, 36, 38, 45[30] B. L. Edelman, S. Goel, S. Kakade, and C. Zhang. Inductive biases andvariable creation in self-attention mechanisms. In K. Chaudhuri, S. Jegelka,L. Song, C. Szepesvari, G. Niu, and S. Sabato, editors, Proceedings of the39th International Conference on Machine Learning, volume 162 ofProceedings of Machine Learning Research, pages 5793–5831. PMLR,17–23 Jul 2022. URL https://proceedings.mlr.press/v162/edelman22a.html.→ page 3356[31] R. Entezari, H. Sedghi, O. Saukh, and B. Neyshabur. The role ofpermutation invariance in linear mode connectivity of neural networks.ArXiv, abs/2110.06296, 2022. → page 18[32] M. Fereydounian, H. Hassani, and A. Karbasi. What functions can graphneural networks generate? 2022. → pages 24, 29[33] C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fastadaptation of deep networks. In D. Precup and Y. W. Teh, editors,Proceedings of the 34th International Conference on Machine Learning,volume 70 of Proceedings of Machine Learning Research, pages 1126–1135.PMLR, 06–11 Aug 2017. URLhttps://proceedings.mlr.press/v70/finn17a.html. → page 52[34] W. Foundation. Wikimedia downloads. URL https://dumps.wikimedia.org.→ page 90[35] S. Garg, D. Tsipras, P. S. Liang, and G. Valiant. What can transformers learnin-context? a case study of simple function classes. In S. Koyejo,S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors,Advances in Neural Information Processing Systems, volume 35, pages30583–30598. Curran Associates, Inc., 2022. URLhttps://proceedings.neurips.cc/paper_files/paper/2022/file/c529dba08a146ea8d6cf715ae8930cbe-Paper-Conference.pdf. → page 47[36] J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl. Neuralmessage passing for quantum chemistry. ArXiv, abs/1704.01212, 2017. →page 12[37] L. S. Gottfredson. Mainstream science on intelligence: An editorial with 52signatories, history, and bibliography, 1997. → page 1[38] A. Graves, G. Wayne, and I. Danihelka. Neural turing machines. ArXiv,abs/1410.5401, 2014. → pages 7, 28[39] I. Gulrajani and D. Lopez-Paz. In search of lost domain generalization.ArXiv, abs/2007.01434, 2021. → page 15[40] N. Harvey, C. Liaw, and A. Mehrabian. Nearly-tight VC-dimension boundsfor piecewise linear neural networks. In S. Kale and O. Shamir, editors,Proceedings of the 2017 Conference on Learning Theory, volume 65 ofProceedings of Machine Learning Research, pages 1064–1068. PMLR,07–10 Jul 2017. URL https://proceedings.mlr.press/v65/harvey17a.html. →page 857[41] H. He. Tweet, 2022. URLhttps://twitter.com/cHHillee/status/1635790330854526981. → page 2[42] D. Hendrycks, N. Mu, E. D. Cubuk, B. Zoph, J. Gilmer, andB. Lakshminarayanan. Augmix: A simple data processing method toimprove robustness and uncertainty, 2020. → page 6[43] S. Hochreiter and J. Schmidhuber. Long short-term memory. NeuralComputation, 9(8):1735–1780, 1997. → page 3[44] B. Ibarz, V. Kurin, G. Papamakarios, K. Nikiforou, M. A. Bennani,R. Csordás, A. Dudzik, M. Bovsnjak, A. Vitvitskyi, Y. Rubanova, A. Deac,B. Bevilacqua, Y. Ganin, C. Blundell, and P. Velicˇkovic´. A generalist neuralalgorithmic learner. ArXiv, abs/2209.11142, 2022. → page 28[45] E. Jang, S. S. Gu, and B. Poole. Categorical reparameterization withgumbel-softmax. ArXiv, abs/1611.01144, 2017. → page 71[46] S. Jelassi, M. E. Sander, and Y.-F. Li. Vision transformers provably learnspatial structure. ArXiv, abs/2210.09221, 2022. → page 33[47] J. Johnson, B. Hariharan, L. van der Maaten, L. Fei-Fei, C. L. Zitnick, andR. Girshick. Clevr: A diagnostic dataset for compositional language andelementary visual reasoning, 2016. → page 1[48] C. Joshi. Transformers are graph neural networks. The Gradient, 2020. →page 4[49] C. K. Joshi, Q. Cappart, L.-M. Rousseau, and T. Laurent. Learning thetravelling salesperson problem requires rethinking generalization.Constraints An Int. J., 27:70–98, 2022. → page 28[50] L. Kaiser and I. Sutskever. Neural gpus learn algorithms. arXiv: Learning,2016. → page 28[51] N. Karalias and A. Loukas. Erdos goes neural: an unsupervised learningframework for combinatorial optimization on graphs. ArXiv,abs/2006.10643, 2020. → page 28[52] J. Kim, T. D. Nguyen, S. Min, S. Cho, M. Lee, H. Lee, and S. Hong. Puretransformers are powerful graph learners. ArXiv, abs/2207.02505, 2022. →pages 24, 6958[53] J. Kim, M. Kim, and B. Mozafari. Provable memorization capacity oftransformers. In The Eleventh International Conference on LearningRepresentations, 2023. URLhttps://openreview.net/forum?id=8JCg5xJCTPR. → pages 34, 36[54] T. N. Kipf and M. Welling. Semi-supervised classification with graphconvolutional networks, 2017. → page 3[55] T. Kojima, S. S. Gu, M. Reid, Y. Matsuo, and Y. Iwasawa. Large languagemodels are zero-shot reasoners. ArXiv, abs/2205.11916, 2022. → page 29[56] N. Lee, K. Sreenivasan, J. D. Lee, K. Lee, and D. Papailiopoulos. Teachingarithmetic to small transformers, 2023. → page 7[57] A. Lewkowycz, A. J. Andreassen, D. Dohan, E. Dyer, H. Michalewski, V. V.Ramasesh, A. Slone, C. Anil, I. Schlag, T. Gutman-Solo, Y. Wu,B. Neyshabur, G. Gur-Ari, and V. Misra. Solving quantitative reasoningproblems with language models. ArXiv, abs/2206.14858, 2022. → page 29[58] H. Li, M. Wang, S. Liu, and P.-Y. Chen. A theoretical understanding ofshallow vision transformers: Learning, generalization, and samplecomplexity. ArXiv, abs/2302.06015, 2023. → page 33[59] Y. Li, F. Gimeno, P. Kohli, and O. Vinyals. Strong generalization andefficiency in neural programs. ArXiv, abs/2007.03629, 2020. → page 69[60] Y. Li, M. E. Ildiz, D. Papailiopoulos, and S. Oymak. Transformers asalgorithms: Generalization and implicit model selection in in-contextlearning. arXiv preprint arXiv:2301.07067, 2023. → page 33[61] R. Liao, R. Urtasun, and R. Zemel. A pac-bayesian approach togeneralization bounds for graph neural networks, 2020. → pages 6, 7[62] S. Luo, S. Li, S. Zheng, T.-Y. Liu, L. Wang, and D. He. Your transformermay not be as powerful as you expect. In S. Koyejo, S. Mohamed,A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in NeuralInformation Processing Systems, volume 35, pages 4301–4315. CurranAssociates, Inc., 2022. URLhttps://proceedings.neurips.cc/paper_files/paper/2022/file/1ba5f64159d67775a251cf9ce386a2b9-Paper-Conference.pdf. → page 33[63] H. Maron, H. Ben-Hamu, H. Serviansky, and Y. Lipman. Provably powerfulgraph networks. ArXiv, abs/1905.11136, 2019. → page 2959[64] S. Maskey, R. Levie, Y. Lee, and G. Kutyniok. Generalization analysis ofmessage passing neural networks on large random graphs. 2022. → page 15[65] J. Miller, R. Taori, A. Raghunathan, S. Sagawa, P. W. Koh, V. Shankar,P. Liang, Y. Carmon, and L. Schmidt. Accuracy on the line: on the strongcorrelation between out-of-distribution and in-distribution generalization.ArXiv, abs/2107.04649, 2021. → pages xi, 14, 16[66] C. Morris and P. Mutzel. Towards a practical k-dimensionalweisfeiler-leman algorithm. ArXiv, abs/1904.01543, 2019. → page 24[67] C. Morris, M. Ritzert, M. Fey, W. L. Hamilton, J. E. Lenssen, G. Rattan, andM. Grohe. Weisfeiler and leman go neural: Higher-order graph neuralnetworks. Proceedings of the AAAI Conference on Artificial Intelligence, 33(01):4602–4609, Jul. 2019. doi:10.1609/aaai.v33i01.33014602. URLhttps://ojs.aaai.org/index.php/AAAI/article/view/4384. → pages 11, 24, 29[68] C. Morris, Y. Lipman, H. Maron, B. Rieck, N. M. Kriege, M. Grohe, M. Fey,and K. M. Borgwardt. Weisfeiler and leman go machine learning: The storyso far. ArXiv, abs/2112.09992, 2021. → page 29[69] B. Neyshabur, S. Bhojanapalli, and N. Srebro. A pac-bayesian approach tospectrally-normalized margin bounds for neural networks, 2018. → page 7[70] OpenAI. Gpt-4 technical report, 2023. → pages 1, 6[71] Y. Ovadia, E. Fertig, J. Ren, Z. Nado, D. Sculley, S. Nowozin, J. Dillon,B. Lakshminarayanan, and J. Snoek. Can you trust your model’suncertainty? evaluating predictive uncertainty under dataset shift. Advancesin neural information processing systems, 32, 2019. → page 14[72] S. Park, J. Lee, C. Yun, and J. Shin. Provable memorization via deep neuralnetworks using sub-linear parameters. In M. Belkin and S. Kpotufe, editors,Proceedings of Thirty Fourth Conference on Learning Theory, volume 134of Proceedings of Machine Learning Research, pages 3627–3661. PMLR,15–19 Aug 2021. URL https://proceedings.mlr.press/v134/park21a.html. →pages 8, 32, 34, 36[73] C. Qian, G. Rattan, F. Geerts, C. Morris, and M. Niepert. Ordered subgraphaggregation networks. ArXiv, abs/2206.11168, 2022. → page 29[74] A. Radford and K. Narasimhan. Improving language understanding bygenerative pre-training. 2018. → page 560[75] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever.Language models are unsupervised multitask learners. 2019. → pages 6, 90[76] A. Radford, J. W. Kim, T. Xu, G. Brockman, C. McLeavey, and I. Sutskever.Robust speech recognition via large-scale weak supervision, 2022. → pages1, 4[77] S. Rajput, K. Sreenivasan, D. Papailiopoulos, and A. Karbasi. Anexponential improvement on the memorization capacity of deep thresholdnetworks. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W.Vaughan, editors, Advances in Neural Information Processing Systems,volume 34, pages 12674–12685. Curran Associates, Inc., 2021. → page 34[78] Y. Razeghi, R. L. L. I. au2, M. Gardner, and S. Singh. Impact of pretrainingterm frequencies on few-shot reasoning, 2022. → page 1[79] Y. Rong, W. Huang, T. Xu, and J. Huang. Dropedge: Towards deep graphconvolutional networks on node classification. In ICLR, 2020. → pagesxi, 15, 16[80] A. Sahiner, T. Ergen, B. Ozturkler, J. Pauly, M. Mardani, and M. Pilanci.Unraveling attention via convex duality: Analysis and interpretations ofvision transformers. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvari,G. Niu, and S. Sabato, editors, Proceedings of the 39th InternationalConference on Machine Learning, volume 162 of Proceedings of MachineLearning Research, pages 19050–19088. PMLR, 17–23 Jul 2022. URLhttps://proceedings.mlr.press/v162/sahiner22a.html. → page 33[81] K. Singhal, T. Tu, J. Gottweis, R. Sayres, E. Wulczyn, L. Hou, K. Clark,S. Pfohl, H. Cole-Lewis, D. Neal, M. Schaekermann, A. Wang, M. Amin,S. Lachgar, P. Mansfield, S. Prakash, B. Green, E. Dominowska, B. A.y Arcas, N. Tomasev, Y. Liu, R. Wong, C. Semturs, S. S. Mahdavi, J. Barral,D. Webster, G. S. Corrado, Y. Matias, S. Azizi, A. Karthikesalingam, andV. Natarajan. Towards expert-level medical question answering with largelanguage models, 2023. → page 1[82] E. D. Sontag. Shattering All Sets of ‘k’ Points in “General Position”Requires (k — 1)/2 Parameters. Neural Computation, 9(2):337–348, 021997. ISSN 0899-7667. doi:10.1162/neco.1997.9.2.337. URLhttps://doi.org/10.1162/neco.1997.9.2.337. → pages 7, 8, 34[83] A. Talmor, J. Herzig, N. Lourie, and J. Berant. Commonsenseqa: A questionanswering challenge targeting commonsense knowledge, 2019. → page 161[84] H. Tang, Z. Huang, J.-Y. Gu, B.-L. Lu, and H. Su. Towards scale-invariantgraph-related problem solving by iterative homogeneous graph neuralnetworks. ArXiv, abs/2010.13547, 2020. → page 16[85] H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix,B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin,E. Grave, and G. Lample. Llama: Open and efficient foundation languagemodels, 2023. → page 31[86] A. van den Oord, O. Vinyals, and K. Kavukcuoglu. Neural discreterepresentation learning. In NIPS, 2017. → page 71[87] G. Vardi, G. Yehudai, and O. Shamir. On the optimal memorization power ofreLU neural networks. In International Conference on LearningRepresentations, 2022. URL https://openreview.net/forum?id=MkTPtnjeYTV.→ pages 8, 32, 34[88] D. Varis and O. Bojar. Sequence length is a domain: Length-basedoverfitting in transformer models. In EMNLP, 2021. → page 10[89] A. Vaswani, N. M. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez,L. Kaiser, and I. Polosukhin. Attention is all you need. In NIPS, 2017. →pages 1, 3, 4, 9, 23, 43[90] P. Velicˇkovic´ and C. Blundell. Neural algorithmic reasoning. Patterns, 2,2021. → page 10[91] P. Velicˇkovic´, G. Cucurull, A. Casanova, A. Romero, P. Lio’, and Y. Bengio.Graph attention networks. ArXiv, abs/1710.10903, 2018. → page 13[92] P. Velicˇkovic´, L. Buesing, M. Overlan, R. Pascanu, O. Vinyals, andC. Blundell. Pointer graph networks. ArXiv, abs/2006.06380, 2020. →pages 13, 28[93] P. Velicˇkovic´, R. Ying, M. Padovano, R. Hadsell, and C. Blundell. Neuralexecution of graph algorithms. ArXiv, abs/1910.10593, 2020. → page 28[94] P. Velicˇkovic´, A. P. Badia, D. Budden, R. Pascanu, A. Banino,M. Dashevskiy, R. Hadsell, and C. Blundell. The clrs algorithmic reasoningbenchmark. 2021. → pages ix, 1, 10, 11, 13, 14, 17, 19, 23, 28, 66, 68[95] R. Vershynin. Memory capacity of neural networks with threshold andrectified linear unit activations. SIAM J. Math. Data Sci., 2:1004–1033, 2020.URL https://api.semanticscholar.org/CorpusID:226352137. → pages7, 32, 34, 36, 4362[96] J. von Oswald, E. Niklasson, E. Randazzo, J. Sacramento, A. Mordvintsev,A. Zhmoginov, and M. Vladymyrov. Transformers learn in-context bygradient descent. ArXiv, abs/2212.07677, 2022. → page 33[97] X. Wang, J. Wei, D. Schuurmans, Q. Le, E. Chi, and D. Zhou.Self-consistency improves chain of thought reasoning in language models.ArXiv, abs/2203.11171, 2022. → page 29[98] J. Wei, Y. Tay, R. Bommasani, C. Raffel, B. Zoph, S. Borgeaud,D. Yogatama, M. Bosma, D. Zhou, D. Metzler, E. Chi, T. Hashimoto,O. Vinyals, P. Liang, J. Dean, and W. Fedus. Emergent abilities of largelanguage models. ArXiv, abs/2206.07682, 2022. → pages 29, 51[99] J. Wei, X. Wang, D. Schuurmans, M. Bosma, E. Chi, Q. Le, and D. Zhou.Chain of thought prompting elicits reasoning in large language models.ArXiv, abs/2201.11903, 2022. → page 29[100] B. Weisfeiler and A. Leman. The reduction of a graph to canonical form andthe algebra which appears therein. Nauchno-Technicheskaya Informatsia, 2(9):12–16, 1968. → page 29[101] F. Wenzel, A. Dittadi, P. V. Gehler, C.-J. Simon-Gabriel, M. Horn,D. Zietlow, D. Kernert, C. Russell, T. Brox, B. Schiele, B. Schölkopf, andF. Locatello. Assaying out-of-distribution generalization in transfer learning.In A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, editors, Advances inNeural Information Processing Systems, 2022. URLhttps://openreview.net/forum?id=57Ryl7lLD4h. → pages 6, 15[102] T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac,T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen,C. Ma, Y. Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest,and A. M. Rush. Huggingface’s transformers: State-of-the-art naturallanguage processing, 2020. → page 45[103] M. Wortsman, G. Ilharco, S. Y. Gadre, R. Roelofs, R. Gontijo-Lopes, A. S.Morcos, H. Namkoong, A. Farhadi, Y. Carmon, S. Kornblith, andL. Schmidt. Model soups: averaging weights of multiple fine-tuned modelsimproves accuracy without increasing inference time. ArXiv,abs/2203.05482, 2022. → page 18[104] Y. Wu, M. Schuster, Z. Chen, Q. V. Le, M. Norouzi, W. Macherey,M. Krikun, Y. Cao, Q. Gao, K. Macherey, J. Klingner, A. Shah, M. Johnson,X. Liu, Łukasz Kaiser, S. Gouws, Y. Kato, T. Kudo, H. Kazawa, K. Stevens,63G. Kurian, N. Patil, W. Wang, C. Young, J. Smith, J. Riesa, A. Rudnick,O. Vinyals, G. Corrado, M. Hughes, and J. Dean. Google’s neural machinetranslation system: Bridging the gap between human and machinetranslation, 2016. → page 3[105] Z. Wu, L. Qiu, A. Ross, E. Akyürek, B. Chen, B. Wang, N. Kim, J. Andreas,and Y. Kim. Reasoning or reciting? exploring the capabilities and limitationsof language models through counterfactual tasks, 2023. → page 2[106] K. Xu, W. Hu, J. Leskovec, and S. Jegelka. How powerful are graph neuralnetworks? ArXiv, abs/1810.00826, 2019. → page 29[107] Y. Yan, K. Swersky, D. Koutra, P. Ranganathan, and M. Hashemi. Neuralexecution engines: Learning to execute subroutines. ArXiv, abs/2006.08084,2020. → pages 4, 28[108] G. Yehuda, M. Gabel, and A. Schuster. It’s not what machines can learn, it’swhat we cannot teach. In ICML, 2020. → page 28[109] G. Yehudai, E. Fetaya, E. A. Meirom, G. Chechik, and H. Maron. Fromlocal structures to size generalization in graph neural networks. In ICML,2021. → pages 15, 20[110] C. Yun, S. Sra, and A. Jadbabaie. Small relu networks are powerfulmemorizers: a tight analysis of memorization capacity. In H. Wallach,H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett,editors, Advances in Neural Information Processing Systems, volume 32.Curran Associates, Inc., 2019. → pages 7, 32, 34[111] C. Yun, S. Bhojanapalli, A. S. Rawat, S. Reddi, and S. Kumar. Aretransformers universal approximators of sequence-to-sequence functions? InInternational Conference on Learning Representations, 2020. URLhttps://openreview.net/forum?id=ByxRM0Ntvr. → page 33[112] C. Yun, Y.-W. Chang, S. Bhojanapalli, A. S. Rawat, S. Reddi, and S. Kumar.O(n) connections are expressive enough: Universal approximability ofsparse transformers. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan,and H. Lin, editors, Advances in Neural Information Processing Systems,volume 33, pages 13783–13794. Curran Associates, Inc., 2020. URLhttps://proceedings.neurips.cc/paper_files/paper/2020/file/9ed27554c893b5bad850a422c3538c15-Paper.pdf. → page 33[113] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Póczos, R. Salakhutdinov, andA. Smola. Deep sets. In NIPS, 2017. → page 1364[114] C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals. Understandingdeep learning requires rethinking generalization. In InternationalConference on Learning Representations, 2017. URLhttps://openreview.net/forum?id=Sy8gdB9xx. → pages 7, 31, 32[115] C. Zhang, M. Raghu, J. M. Kleinberg, and S. Bengio. Pointer value retrieval:A new benchmark for understanding the limits of neural networkgeneralization. ArXiv, abs/2107.12580, 2021. → pages 10, 28[116] H. Zhang, M. Cissé, Y. Dauphin, and D. Lopez-Paz. mixup: Beyondempirical risk minimization. ArXiv, abs/1710.09412, 2018. → pagesxi, 6, 15, 16[117] Y. Zhang, A. Backurs, S. Bubeck, R. Eldan, S. Gunasekar, and T. Wagner.Unveiling transformers with lego: a synthetic reasoning task. ArXiv,abs/2206.04301, 2022. → page 28[118] H. Zhou, A. Nova, H. Larochelle, A. Courville, B. Neyshabur, and H. Sedghi.Teaching algorithmic reasoning via in-context learning, 2022. → page 7[119] Y. Zhou, G. Kutyniok, and B. Ribeiro. Ood link prediction generalizationcapabilities of message-passing gnns in larger test graphs. ArXiv,abs/2205.15117, 2022. → page 15[120] Łukasz Kaiser and I. Sutskever. Neural gpus learn algorithms, 2016. → page365Appendix AAppendix for OOD GeneralizationA.1 Experiment DetailsFor Table 2.1 we take the (with-hint versions) numbers from [94] and run the no-hintversion of the table with the same configuration except for disabling hints.For the rest of the experiments, we use a batch size of 32 for message-passingprocessors and 16 for high-order processors, gradient clipping to a norm 1.0, alearning rate of 0.0001, 20000 optimization steps, and cosine annealing scheduler.Moreover, as the number of processing steps is meaningless in the no-hint version,we set this number to 32 (i.e. 32 recurrent steps of message passing). Unlessotherwise specified, experiments are run using the MPNN-G processor on theL-CLRS dataset. Finally, as the ID accuracy and OOD accuracy are not highlycorrelated, we do not do any early-stopping and among the models with the highestvalidation accuracy, we keep the last model.The number of distinct algorithms shrinks from 30 to 24 when doing suchchanges. In particular, the following algorithms merge together:• Sorting: All variants of sorting. That is, Quick sort, Insertion sort, Bubblesort, Heap sort are merged together. We Keep the “Quick sort” name torepresent sorting and refrain from creating new names.• Single source shortest path: Dijkstra and Bellman-Ford are merged together.We keep the “Bellman Ford” name to represent the shortest path.66• String Matching: Naive String Matcher and KMP algorithms are mergedtogether as they contain the same input/output. We keep the “Naive StringMatcher” name to represent the string matching task.• Convex Hull computation: Jarvis March and Graham Scan are combined andrepresented as “Graham Scan”.Note that some other algorithms might also seem to have the same computation(e.g., MST Prim and MST Kruskal both compute minimum spanning tree), however,they are implemented differently in the CLRS benchmark and we keep them differentas well.A.2 Leveraging Intermediate Steps of AlgorithmsThe CLRS benchmark contains intermediate steps of algorithms (i.e., hints) as op-tional supervision in addition to the supervision of the final output of the algorithm.Namely, for applying an algorithm to a particular graph, two additional supervisionsignals are provided: (1) number T representing the number of recurrent processingsteps to perform (i.e., the processor is called T times), (2) depending on the algo-rithm, several hints are provided. Each hint is a T -dimensional tensor, representingthe intermediate output of an algorithm for each intermediate step. For instance,for the Dijkstra algorithm, one hint is the computed distance for each node at eachindividual step (i.e., Xhint ∈ RT×n where Xhinttv represents the distance of node v tothe source node at step t). Hints also follow the encode-process-decode frameworkand have separate encoders and decoders (while sharing the same processor withthe inputs and outputs). When leveraging hints, Eq. (2.1) should be modified to thefollowing formula:H(t+1) = GNN([ H(t), Zinpv , Z(t)hintv],A), (A.1)where Z(t)hintv is the encoded hints at step t (1≤ t ≤ T ). The hints for the successorstep are also decoded and evaluated in the final loss function along with the finaloutput of the algorithm.67Table A.2 contains the full OOD test accuracies of the algorithms when usinghints and when disabling them. As shown by the results, enabling hints mostlydegrade the final OOD accuracy. While [94] employed a complex model selectionon the validation numbers for the with-hint version of processors, running a no-hintversion of the simplest form of a processor and no hyperparameter tuning yieldssuperior OOD results.A.3 Encode-Process-Decode FrameworkHere, we explain how three algorithms are put into a simplified encode-process-decode framework.Sorting. Given a sequence a1,a2, . . . ,an of scalars, the task is to find the successor ofeach element when the sequence is sorted. So, the input is again a two-dimensionalnode feature matrix XV ∈ Rn×2 where the first dimension denotes the index of eachelement and the second dimension is the value ai of that dimension. Here, we donot have an underlying graph and one could consider a fully-connected graph whereall nodes are connected to each other. The encoding, processing, and decoding stepsare similar to the case of BFS since the final output is again pointer prediction foreach node.Minimum Spanning Tree (MST) Prim. Given a weighted graph, the goal is tofind the minimum spanning tree of the input graph. A spanning tree is a subsetof edges with the lowest sum of weights, which preserves the connectivity of thegraph. The input to this task is a two-dimensional node feature matrix XV ∈ Rn×2where the first dimension is the node index, and the second dimension indicatedthe starting node of the MST Prim algorithm. Additionally, the input contains anadjacency matrix A ∈ Rn×n and edges features/weights XE ∈ Rn×n×1. Again, theprocedure is similar to that of BFS, but with the following differences: during theencoding part, we also need an edge encoder (e.g., WE ∈ R1×dh ) to encode edgefeatures and incorporate them in the processing step. Finally, the decoder needsto predict a binary output for each edge (i.e., whether the edge is in MST or not).So, for each edge ei j with two endpoint nodes i, j, respectively, we directly applya sigmoid function to the output of the decoder αi j, and compute the binary crossentropy loss with the ground truth edge label.68A.4 Alternative Representation for Input IndexRepresentation ShiftThe advantage of the Random Scalar Index is that it is a simple extension of theposition index into node embeddings. However, it requires a large number oftraining steps for the neural network to learn the ordering of numbers in the range[0, 1]. Ideally, a representation method for ordering of nodes must be invariantunder shifting and scaling of the index (i.e., if we shift the index of nodes by someconstant factor, the function must not change).Edge Position. A more natural alternative way of enforcing such an embeddingis to encode orderings as edge features [59]. Namely, we create an edge featurefrom node vi to node v j as follows:posvi,v j = 1 [i< j] , ∀i, j ∈ {1,2, . . . ,n}.This way, the ordering of nodes can be determined by performing a topological sortusing two steps of message passing on a fully connected graph. The advantage ofusing such a method is that nodes do not undergo any distribution shift as the graphgrows; hence any subgraph of the larger graph stays in distribution. Moreover, themethod is invariant to any shift or scaling of indices i, j. On the other hand, usingthis method requires a processor with global message passing steps and requiresan implicit extrapolation of the Topological Sort itself (which is itself challenging,according to the Topological Sort task in the benchmark).In order to keep the expressiveness of the models and preserve the uniquenessof positions on each node, we append a vector of random features to input featuresof each node for all representation strategies [52]. Since this method requires afully-connected MPNN, we run this method on the MPNN-FC processor and reportthe results in Table A.1. The results show that this method is indeed superior toScalar Index, although it is outperformed by Random Scalar by a narrow margin.A.5 Additional Approaches that FailedFurthermore, we also tried the following ideas (mostly on the DFS algorithm, andmostly towards making hints work better than no-hints), however, they seem to not69Table A.1: Performance of different positional encoding strategies onpositions-order aware algorithms for MPNN-FC on L-CLRS dataset.Method Scalar Random Scalar Edge Pos.Dfs 15.07± 0.92 25.29± 1.25 28.19± 1.20Find Maximum Subarray Kadane 23.81± 2.12 48.99± 2.78 46.27± 5.29Dag Shortest Paths 99.67± 0.12 99.72± 0.06 99.71± 0.04Naive String Matcher 11.51± 4.26 7.13± 3.35 6.48± 3.43Matrix Chain Order 65.19± 3.05 91.50± 0.59 80.92± 1.71Topological Sort 68.77± 2.61 70.06± 0.53 67.25± 0.90Bfs 99.90± 0.08 100.00± 0.00 99.82± 0.10Mean 54.85 63.24 61.23work:• Markovian Assumption on hints: One key factor in failure in OOD general-ization is a large number of roll-outs mostly present in sequential algorithms).One scenario where hints might be useful is to consider a Markovian assump-tion on steps of the algorithm, meaning that each step of the algorithm isindependent of previous steps given its previous state. This then allows forstep-wise training and makes training easier and less memory/compute hun-gry. However, this strategy (on algorithms that satisfy Markovian property)did not bring any performance gain over non-markovian assumption.• Choosing a subset of hints: Each algorithm might contain up to 22 differenthints with different types. However, not all of them have a positive effect on aGNN for a better prediction of output and might distract the GNN more thanthey can help. We tried using only a subset of such hints to see if any boost inoutput happens or not. For the DFS algorithm, we observed that the fewerhints we use, the higher the accuracy becomes, however, it never surpassesthe no-hint version.• Using hints only as masks: Since an algorithm can contain several hints andnot all of them are useful to the GNN, we hypothesized that perhaps a GNNshould have fewer restrictions. We used the change in hints as a mask to70supervise the GNN for a sparser hidden-representation update. However,again it did not bring any performance gain over the no-hint counterpart andthe GNN (on the DFS algorithm) violated the mask quite often (in favor ofalways updating hidden representations). Even increasing the mask loss toprevent such violations, reduces learning accuracy.• Discretize/Quantize hidden representations: Due to the long roll-out of manyalgorithms, small errors might accumulate throughout the processing stepsand lead to wrong outputs. We tried quantization techniques such as StraightThrough [9], Gumbel-Softmax [45], Vector Quantization [86] on the hiddenrepresentation of each processing step. However, training deep quantizednetworks in an end-to-end fashion hurts the in-distribution training accuracyand even in cases where it did not, OOD accuracy was no better than thecontinuous counterpart. Even when using hints and trying a Markovianassumption, quantization did not offer any accuracy gain.71Table A.2: Full test results of Table 2.1.Method MPNN-FC PGN GATUse Hints Yes No Yes No Yes NoArticulation Points 50.91± 2.18 50.36± 3.02 49.53± 2.09 47.85± 2.59 37.76± 1.62 49.24± 1.30Activity Selector 80.66± 3.16 91.49± 1.29 66.80± 1.62 64.59± 1.26 73.23± 1.37 66.14± 0.63Bellman Ford 92.01± 0.28 90.14± 1.17 92.99± 0.34 91.59± 1.99 87.91± 1.19 78.12± 0.24Bfs 99.89± 0.05 99.53± 0.20 99.63± 0.29 99.79± 0.20 99.04± 0.21 99.14± 0.68Binary Search 36.83± 0.26 7.06± 6.87 76.95± 0.13 29.13± 19.12 23.50± 3.12 3.58± 0.78Bridges 72.69± 4.78 30.36± 0.73 51.42± 7.82 32.90± 0.99 25.64± 6.60 31.52± 0.49Bubble Sort 5.27± 0.60 50.13± 3.33 6.01± 1.95 54.56± 4.23 9.91± 1.77 57.28± 1.55Dag Shortest Paths 96.24± 0.56 97.23± 0.17 96.94± 0.16 97.46± 0.35 81.14± 1.37 80.60± 1.48Dfs 6.54± 0.51 15.27± 5.29 8.71± 0.24 22.64± 3.12 11.78± 2.04 11.13± 1.31Dijkstra 91.50± 0.50 93.51± 0.46 83.45± 1.75 93.80± 0.22 58.01± 0.79 83.27± 0.17Find Maximum Subarray Kadane 20.30± 0.49 21.65± 0.58 65.23± 2.56 15.45± 0.31 24.43± 0.43 14.88± 0.52Floyd Warshall 26.74± 1.77 39.47± 4.43 28.76± 0.51 35.19± 1.20 16.66± 3.14 14.78± 3.85Graham Scan 91.04± 0.31 96.27± 0.27 56.87± 1.61 57.05± 0.29 77.89± 2.70 57.62± 0.33Heapsort 10.94± 0.84 63.72± 6.41 5.27± 0.18 55.18± 2.05 10.35± 1.83 59.55± 1.64Insertion Sort 19.81± 2.08 38.41± 18.19 44.37± 2.43 68.77± 2.56 29.52± 1.87 40.93± 18.49Jarvis March 34.86± 12.39 95.46± 1.01 49.19± 1.07 56.64± 0.55 51.51± 10.25 57.63± 0.33Kmp Matcher 2.49± 0.86 9.08± 2.39 2.00± 0.12 2.18± 0.53 3.03± 0.36 2.28± 0.27Lcs Length 53.23± 0.36 56.93± 2.87 56.82± 0.21 60.27± 0.35 57.88± 1.02 59.70± 0.68Matrix Chain Order 79.84± 1.40 78.13± 6.56 83.91± 0.49 82.01± 0.59 78.19± 3.31 75.13± 5.57Minimum 85.34± 0.88 99.53± 0.06 87.71± 0.52 98.01± 1.13 84.20± 2.95 93.52± 4.88Mst Kruskal 70.97± 1.50 73.22± 0.39 66.96± 1.36 71.25± 1.48 65.72± 0.99 71.57± 2.81Mst Prim 69.08± 7.56 74.63± 3.19 63.33± 0.98 73.16± 1.08 38.20± 4.34 45.46± 4.03Naive String Matcher 3.92± 0.30 4.43± 0.84 2.08± 0.20 1.94± 0.27 3.01± 1.20 1.30± 0.20Optimal Bst 62.23± 0.44 56.02± 13.13 71.01± 1.82 30.25± 30.03 65.49± 1.75 30.42± 2.85Quickselect 1.43± 0.69 8.50± 3.29 3.66± 0.42 9.10± 1.07 4.36± 0.95 8.69± 0.99Quicksort 11.30± 0.10 58.30± 4.20 6.17± 0.15 63.87± 3.06 7.60± 0.98 59.29± 0.92Segments Intersect 93.44± 0.10 93.94± 0.29 77.51± 0.75 77.67± 0.92 90.41± 0.04 76.16± 0.73Strongly Connected Components 24.37± 4.88 53.65± 3.52 20.80± 0.64 57.65± 3.95 12.70± 3.12 39.86± 2.19Task Scheduling 84.11± 0.32 82.86± 0.32 84.89± 0.91 81.88± 0.11 84.69± 2.09 81.93± 0.14Topological Sort 52.60± 6.24 72.03± 0.74 60.45± 2.69 78.76± 2.69 27.03± 6.92 35.30± 19.58Mean 51.02±0.60 60.04±0.93 52.31±0.35 57.02±1.23 44.69±0.58 49.53±0.97Wins 8/30 22/30 10/30 20/30 14/30 16/30Validation Score Mean 96.63 93.83 89.47 83.79 95.66 78.8372Table A.3: Validation scores of Table 2.5.Method MPNN-G 2WL Average Hybrid Sigmoid HybridArticulation Points 100.00± 0.00 100.00± 0.00 100.00± 0.00 100.00± 0.00Activity Selector 99.49± 0.18 99.49± 0.48 99.23± 0.31 99.10± 0.18Bellman Ford 100.00± 0.00 100.00± 0.00 100.00± 0.00 100.00± 0.00Bfs 100.00± 0.00 100.00± 0.00 100.00± 0.00 100.00± 0.00Binary Search 99.10± 0.08 99.41± 0.11 99.45± 0.15 99.09± 0.22Bridges 99.08± 0.08 100.00± 0.00 100.00± 0.00 100.00± 0.00Dag Shortest Paths 100.00± 0.00 99.80± 0.28 100.00± 0.00 100.00± 0.00Dfs 99.93± 0.09 99.67± 0.18 99.22± 0.55 99.80± 0.16Find Maximum Subarray Kadane 97.51± 0.21 97.46± 0.67 97.46± 0.18 97.48± 0.05Floyd Warshall 88.77± 0.90 87.37± 0.59 88.46± 0.63 87.46± 0.42Graham Scan 100.00± 0.00 100.00± 0.00 100.00± 0.00 100.00± 0.00Lcs Length 75.00± 0.35 100.00± 0.00 99.98± 0.02 100.00± 0.00Matrix Chain Order 99.52± 0.03 99.33± 0.06 99.49± 0.12 99.45± 0.10Minimum 99.97± 0.02 99.97± 0.05 99.98± 0.02 99.97± 0.05Mst Kruskal 92.90± 0.08 99.19± 0.33 99.77± 0.17 99.88± 0.17Mst Prim 99.35± 0.18 99.35± 0.09 99.61± 0.16 98.96± 0.18Naive String Matcher 80.89± 14.72 78.45± 10.80 83.15± 4.53 92.01± 4.73Optimal Bst 96.23± 0.13 95.77± 0.17 96.30± 0.08 95.71± 0.51Quickselect 89.19± 5.67 99.82± 0.02 99.80± 0.07 99.67± 0.13Quicksort 100.00± 0.00 100.00± 0.00 100.00± 0.00 100.00± 0.00Segments Intersect 98.49± 0.10 96.42± 0.53 97.23± 0.47 96.86± 0.67Strongly Connected Components 100.00± 0.00 99.28± 0.49 100.00± 0.00 100.00± 0.00Task Scheduling 99.50± 0.07 99.90± 0.14 99.95± 0.07 99.85± 0.12Topological Sort 100.00± 0.00 100.00± 0.00 100.00± 0.00 99.97± 0.05Mean 96.46 97.95 98.30 98.55Table A.4: Validation scores of Table 2.2.Method CLRS-Len-Deg CLRS-Len CLRS-DegBellman Ford 99.81± 0.02 99.82± 0.02 99.82± 0.02Bfs 100.00± 0.00 100.00± 0.00 100.00± 0.00Dag Shortest Paths 99.98± 0.01 99.98± 0.01 99.98± 0.01Dfs 100.00± 0.00 100.00± 0.00 99.63± 0.06Floyd Warshall 87.03± 0.52 86.58± 0.40 78.40± 0.98Mst Kruskal 62.46± 1.09 62.22± 0.92 63.41± 0.87Mst Prim 98.56± 0.07 98.43± 0.21 97.36± 0.44Strongly Connected Components 99.86± 0.05 99.83± 0.01 99.66± 0.07Topological Sort 99.94± 0.03 99.93± 0.03 99.69± 0.06Mean 94.18 94.09 93.1173Table A.5: Test scores of proposed processors on 12 representative algorithmson (low-data regime) CLRS dataset. While the 2WL processor showssome level of robustness to the low-data regime, MPNN outperforms2WL in shortest-path-related tasks since max aggregation inductive biasis baked into the architecture, while the 2WL architecture requires ahigher amount of data to learn such inductive biases. However, the hybridprocessor still generalizes better compared to both processors.Method MPNN-G 2WL Average HybridArticulation Points 49.78± 0.64 48.66± 6.92 51.74± 3.83Bellman Ford 96.16± 0.88 92.55± 1.50 91.02± 3.16Bfs 100.00± 0.00 99.67± 0.17 100.00± 0.00Dag Shortest Paths 98.45± 0.20 88.64± 4.73 96.22± 1.61Dfs 26.19± 1.20 25.29± 9.87 27.08± 1.54Floyd Warshall 25.92± 2.37 21.74± 6.25 28.94± 1.23Graham Scan 96.81± 0.79 73.80± 3.69 96.88± 0.35Lcs Length 56.60± 1.71 86.58± 0.43 87.54± 0.39Matrix Chain Order 94.81± 0.76 83.66± 5.27 95.24± 1.29Mst Prim 64.34± 6.82 53.47± 0.60 59.70± 2.69Quicksort 48.73± 8.64 42.77± 15.35 50.99± 28.89Task Scheduling 83.10± 0.21 82.47± 0.41 83.10± 0.21Mean 70.07 66.61 72.37Table A.6: Test scores of different hybrid models for combinations of MPNN-G and 2WL on L-CLRS dataset. A hybrid model of two different proces-sors brings accuracy gain over each individual processor.Method Average (MP-MP) Avergae (2WL-2WL) Avergae (MP-2WL) Sigmoid (MP-2WL)Mean 68.77±0.61 69.57±0.64 70.92±0.69 71.98±1.9474Table A.7: Full test scores of Table 2.5 for Scalar Index.Method MPNN-G 2WL Average Hybrid Sigmoid HybridArticulation Points 88.92± 3.39 69.59± 5.05 78.58± 1.48 80.70± 2.11Activity Selector 95.81± 0.67 90.15± 1.03 90.43± 2.65 89.37± 2.31Bellman Ford 98.83± 0.18 98.54± 0.60 98.19± 0.14 98.03± 0.68Bfs 100.00± 0.00 100.00± 0.00 100.00± 0.00 100.00± 0.00Binary Search 93.20± 2.07 90.30± 1.32 86.17± 11.38 89.70± 2.80Bridges 89.17± 2.13 85.22± 11.89 98.69± 0.51 82.57± 17.44Dag Shortest Paths 99.72± 0.06 95.91± 1.54 99.10± 0.14 99.67± 0.06Dfs 25.03± 1.62 21.11± 4.18 19.04± 1.79 22.56± 1.08Find Maximum Subarray Kadane 26.48± 1.50 44.03± 0.53 46.01± 2.10 43.68± 4.64Floyd Warshall 35.39± 2.05 30.75± 0.27 34.38± 0.98 33.59± 1.36Graham Scan 98.88± 0.25 82.69± 9.57 95.26± 1.30 92.93± 1.87Lcs Length 58.85± 0.80 87.06± 0.71 86.97± 0.54 88.56± 0.32Matrix Chain Order 73.65± 9.55 76.29± 3.82 73.47± 5.28 68.70± 3.11Minimum 99.56± 0.04 99.15± 0.22 99.46± 0.11 99.58± 0.20Mst Kruskal 72.07± 1.03 84.77± 2.24 90.64± 0.79 88.01± 2.31Mst Prim 83.84± 1.47 75.05± 2.91 78.37± 1.04 78.03± 2.18Naive String Matcher 7.36± 1.64 8.48± 1.29 9.54± 2.72 12.83± 1.30Optimal Bst 61.43± 11.78 75.94± 0.82 74.96± 1.87 76.12± 2.06Quickselect 3.74± 5.29 0.20± 0.21 0.23± 0.18 29.36± 41.52Quicksort 10.06± 3.49 17.25± 12.53 10.86± 5.96 14.45± 7.63Segments Intersect 98.16± 0.20 96.58± 0.64 97.51± 0.45 96.37± 0.49Strongly Connected Components 75.15± 0.33 55.99± 1.79 72.15± 1.42 73.83± 3.12Task Scheduling 83.64± 0.14 83.85± 0.55 84.21± 0.72 83.82± 0.24Topological Sort 82.61± 1.50 62.09± 8.14 77.86± 3.26 85.00± 2.53Mean 69.23±0.97 67.96±0.64 70.92±0.73 71.98±1.9475Table A.8: Full test scores of Table 2.5 for Random Scalar Index. TheWeighted Mean is average over all 30 algorithms when considering re-peated algorithms with higher weight.Method MPNN-G 2WL Average Hybrid Sigmoid HybridArticulation Points 89.86± 1.79 72.46± 2.51 79.07± 0.88 81.97± 2.99Activity Selector 94.02± 1.96 89.94± 1.02 92.31± 1.40 91.07± 1.26Bellman Ford 98.84± 0.25 98.58± 0.38 98.16± 0.15 96.74± 2.57Bfs 100.00± 0.00 100.00± 0.00 100.00± 0.00 100.00± 0.00Binary Search 95.41± 1.35 94.94± 0.96 94.45± 0.90 96.70± 2.15Bridges 77.95± 8.57 77.96± 9.75 97.05± 2.08 96.43± 0.82Dag Shortest Paths 99.76± 0.07 96.60± 1.52 99.40± 0.14 99.58± 0.20Dfs 30.55± 1.85 37.50± 2.25 26.11± 1.71 25.42± 0.78Find Maximum Subarray Kadane 49.25± 3.93 74.30± 3.57 75.70± 1.99 69.29± 5.46Floyd Warshall 35.04± 1.26 30.32± 0.42 33.75± 0.90 32.66± 1.80Graham Scan 98.41± 0.20 92.27± 2.40 95.80± 0.77 93.00± 1.80Lcs Length 58.96± 0.53 88.40± 1.03 87.53± 0.76 88.99± 0.60Matrix Chain Order 92.34± 0.53 91.59± 1.08 92.54± 2.17 92.12± 1.05Minimum 99.59± 0.19 99.43± 0.12 99.46± 0.16 99.50± 0.08Mst Kruskal 72.78± 0.50 87.56± 2.75 90.42± 1.75 86.69± 2.59Mst Prim 81.79± 1.24 76.64± 2.94 78.22± 1.97 78.52± 1.50Naive String Matcher 5.62± 1.71 14.08± 5.14 16.24± 6.82 14.88± 7.65Optimal Bst 76.96± 1.19 72.49± 2.43 77.68± 1.15 76.82± 2.83Quickselect 3.34± 4.72 0.15± 0.08 0.63± 0.69 0.05± 0.04Quicksort 6.41± 1.98 7.41± 1.56 15.33± 8.84 26.24± 19.85Segments Intersect 97.74± 0.23 96.17± 0.32 97.55± 0.78 97.34± 0.55Strongly Connected Components 75.21± 1.48 50.65± 1.37 71.47± 1.33 72.14± 2.07Task Scheduling 84.04± 0.81 83.92± 0.58 84.12± 0.66 83.69± 0.28Topological Sort 82.80± 0.96 83.18± 1.36 80.84± 1.84 83.65± 1.26Mean 71.11±0.57 71.52±0.53 74.33±0.49 74.31±0.97Weighted Mean 64.29 64.79 68.00 68.89760.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0AccuracyIDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0IDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0IDOOD(a) Task Scheduling0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0AccuracyIDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0IDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0IDOOD(b) Optimal BST0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0AccuracyIDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0IDOOD0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0Interpolation Coefficient0.00.20.40.60.81.0IDOOD(c) MST KruskalFigure A.1: Model interpolation charts on three additional tasks. For each task,we train three models and pairwise interpolate their weights and plot IDand OOD accuracies. Each plot shows the interpolation of one pair foreach task. As suggested by the plots, ID and OOD accuracies have verydifferent landscapes and ID accuracy is a suboptimal model selector. In(a), the right plot shows the highest OOD accuracy corresponds with thelowest ID accuracy. In (b), the middle plot achieves its highest OODaccuracy when ID accuracy is not at its peak. In (c), the right plot showsdifferent rates of change for ID and OOD accuracies, OOD accuracy isat its peak when ID is not, and for coefficient 0.5, the OOD accuracy ishigher than that of ID.77Appendix BAppendix for MemorizationSymbol DescriptionA A single-layer multi-head Attention (MHA) mechanismT Training setH Number of heads in the MHA layerT Training set sizen Context lengthW Kh,W Qh ∈ Rde×dh ,W V h ∈ Rde×dv Weight matrices for each head h ∈ [H]W O ∈ RHdv×de Weight matrix to combine the outputs of different headsW D ∈ Rde Final read-out weight of the modelW The entire set of parameters of the MHA layerE ∈ Rn×de Key/Context matrix containing n key/context vectorse ∈ Rde Query vectory ∈ R Ground truth labelα h,θ h,zh, ph Intermediate variables in the MHA layero Output of the MHA layeryˆ ∈ R Final output of the modelTable B.1: Important notations used in the paper.B.1 Proof of Lemma 1For convenience let us define n1 = min(n,dh). WLOG let us assume that T ≥H(n1−1)+1. This is because for the case of T < H(n1−1)+1 we can add extra78dummy data points that satisfy the assumptions and make T = H(n1−1)+1, thenprove Z has full row-rank, and finally remove the extra rows of Z (i.e. dummy datapoints) to prove the statement.We prove the lemma by induction on H. We will demonstrate that by addingeach head, we can increase the rank of Z by n1−1. The key to proving the inductionis the following two claims:Claim 1 (Fit n1 examples in one head). Assume that n1 examples{(E (t),e(t))}n1t=1satisfy Assumptions 1 and 2. Fix any head index h ∈ [H]. For any arbitrary setof convex combination vectors θ (t)h ∀t ∈ [n1], there exists an Attention head withkey/query weights W Kh,W Qh which achieves such coefficients for the set of n1 datapoints.Claim 2 (Fit n1 − 1 examples in one head, leave the rest intact). Assume thatT ′ ≥ n1 examples{(E (t),e(t))}T ′t=1satisfy Assumptions 1 and 2, and defineR ={T ′−n1+2, . . . ,T ′−1,T ′} as the set of last n1−1 examples andRc = [T ]−R asits complement for convenience. Fix any head index h∈ [H]. Then, for any arbitraryset of desired convex combination vectors p̂(t), t ∈R, there exists key/query weightsW Kh,W Qh (that can depend on p̂(t), t ∈R) such that both θ (t)h = p̂(t),∀t ∈R andθ (t)h , t ∈Rc are not dependent on p̂(t), t ∈R.Claim 1 states that we can design an attention head that produces any arbitraryθ h coefficients for n examples. Furthermore, Claim 2 reinforces Claim 1 by demon-strating the possibility of tuning coefficients θ h of n1−1 examples while keepingthe rest of the examples fixed (and perhaps uncontrollable). However, this comes atthe cost of fitting one less example into a single head. The proofs of both claimswill be deferred until the end of the lemma. We are now ready to prove the lemmaby induction on H.Step 1. The base of induction H = 1. First, we prove the base of induction forH = 1. Let us only consider the first n1 examples (i.e., T ′ = n1). Notice that thematrix Z has n1 rows and d columns in this case. Moreover, the t-th row of Z hasthe formZ t = E (t)⊤θ (t)1 =[e(t)1 ,e(t)2 , . . . ,e(t)n]θ (t)1 ,and is a convex combination of the context tokens for the example t. According to79Assumption 2, the span of context vectors e(t)[n] has dimension n. Therefore, we caninductively find θ (1)1 ,θ(2)1 , . . . ,θ(n1)1 such that each row Z t is not in the span of theprevious t−1 rows. We can continue this process for all t ≤ n1 since the span of theprevious rows has at most t−1 dimensions, but the convex hull of e(t)[n] has at least nlinearly independent elements (n1 ≤ n), so it is not covered by the previous t−1rows. Therefore, there exists {θ (t)1 }n1t=1 such that the resulting Z matrix has rank atleast n. Then, using the result of Claim 1, there exists Attention weights W K1,W Q1that produce the desired set of convex combinations {θ (t)1 }nt=1, and subsequentlythe first n rows of the desired feature matrix Z .Step 2. Induction Step. Next, assume that for H = H ′, there exists heads withAttention weights {(W Kh,W Qh)}H′h=1 such that matrix Z has rank H′(n1− 1)+ 1.We prove for H = H ′+1 there exists Attention weights {(W Kh,W Qh)}H′+1h=1 suchthat the resulting matrix Z attains rank (H ′+ 1)(n1− 1)+ 1. Let us refer to thematrices derived from the induction hypothesis by a prime superscript (i.e., Z ′, and{(W K ′h,W Q′h)}H′h=1). Now, for the case of H = H′+1, we set the weights of the firstH ′ Attention heads to be equal to that of H ′. Namely,W Kh :=W K ′h, W Qh :=W Q′hfor all h ∈ [H ′]. Then, the feature matrix Z can be written as the following blockform:Z = Z ′ Z ′′︸ ︷︷ ︸H ′d︸︷︷︸d ,where Z ′ ∈ RT×dH ′ is the matrix derived from the hypothesis and Z ′′ ∈ RT×d is theblock corresponding to the new head H ′+1, whose parameters we have not yet de-termined. Now, we know rank(Z ′)≥H ′(n1−1)+1. Define k := rank(Z ′). WLOG,assume that the first k rows of the matrix Z ′ are linearly independent (otherwise, wecan rearrange the rows of Z , i.e. the examples, to satisfy this property). Now, takethe last n1−1 rows that are linearly dependent on the first k rows. We defer the casefor when less than n1−1 such rows exist (i.e., k+n1−1> T ) to the later stages of80the proof.Let us partition the matrix Z into two blocks: the first T −n1+1 rows and thelast n1−1 rows. Namely,Z =[ZaZb]= Z′a Z′′aZ ′b Z′′b .Observe that the rows of Z ′b are all linear combinations of rows of Z′a. So, we canzero out the Z ′b block by adding a linear combination of rows of Z′a. Such operationsadd a combination of rows of Z ′′a to the Z′′b as well. Notice that the rank of Z staysthe same since such operations do not change rank. So let us call Z˜ the resultingmatrix:Z˜ = Z′a Z′′a0 Z ′′b +A ,where A ∈ R(n1−1)×d contains a linear combination of rows of Z ′′a . Sincerank(Z)= rank(Z˜)Lemma 2≥ rank(Z ′′b+A)+rank(Z ′a)≥ rank(Z ′′b+A)+H ′(n1−1)+1,we have to find suitable Attention weights such that rank(Z ′′b +A) ≥ n1 − 1 toachieve a (H ′+ 1)(n1− 1)+ 1 rank for Z and complete the proof. Now, Claim2 states that there exists a fixed set of convex combination vectors {θ (t)H ′+1}T−n1t=1such that for any set of convex combination vectors {θ (t)H ′+1}Tt=T−n1+1 there existsAttention weights W KH ′+1,W QH ′+1 that produces such convex combinations. Thismeans that we can assume A is a fixed matrix, and have the flexibility of setting Z ′′bto be any convex combination of context tokens. Similar to the base of induction,we inductively build the individual rows of Z ′′b , but the difference is that the originalmatrix is added to a fixed matrix. Hence, at each inductive step t, the following81constraint must be satisfied:[Z ′′b +A]t /∈ span{[Z ′′b +A]1 ,[Z ′′b +A]2 , . . . ,[Z ′′b +A]t−1},which can be satisfied by a stronger constraint:[Z ′′b]t /∈ span{At ,[Z ′′b +A]1 ,[Z ′′b +A]2 , . . . ,[Z ′′b +A]t−1}.Hence, since the input context tokens are in an n dimensional space, we can continuefor only n−1 rows, find suitable {θ (t)H ′+1}Tt=T−n1+1, and build a full-rank matrix Z ′′b .Thus, applying Claim 2, there exists a set of attention weights W KH ′+1,W QH ′+1that produces rank(Z ′′b +A) ≥ n1− 1. Note that Claim 2 archives the results byletting c → ∞ inside W H ′+1, which is not well-defined. However, to achieverank(Z ′′b +A)≥ n1−1, we only need a finite large enough c since the limit matrixrank(Z ′′b +A) can be arbitrarily approximated, and the rank of a matrix does notdecrease under small perturbations.1 Therefore, this completes the proof.The case of less than n1− 1 linearly dependent rows of Z ′. In this case, wehave k+n1−1 > T . Let us define m := T − k < n1−1. Since we have assumedT ≥ (H ′+1)(n1−1)+1, this means that that Z ′ already has a rank at least (H ′+1)(n1−1)+1−m. So, the head index H ′+1 needs to increase the rank of Z ′ bym. Since m≤ n1−1, all the above steps also apply to m and the rest of the proofprocedure is the same.B.1.1 Proof of Claims 1 and 2.First, note each θ (t)h is derived from applying Softmax nonlinearity on the logitsα (t)h . So, we need to solve a system of equations for α(t)h for each t ∈ [T ]2. Since ourfocus is on a single head, we omit the subscript h for simplicity. Furthermore, wedefine W :=W KW Q⊤. Now, for data indices t1 < t2 < · · ·< tm ∈ [T ] where m≤ n,1This is due to Weyl’s singular value theorem. To ensure a matrix M ∈Rn×d does not lose rank byε perturbation of each entry, one could choose ε < σmin(M)/nd where σmin is the smallest non-zerosingular value of M .2The set of logits that produce a convex combination is not necessarily unique. However, unique-ness is not a requirement in this case, as we only require the existence of some set of logits thatproduces a convex combination.82we are interested in solving the following system of equations:E (ti)W e(ti) = α (ti) ∀i ∈ [m].Since each E (ti) ∈ Rn×d is full-rank with rank n, it has a pseudo right inverseE (ti)† ∈ Rd×n such that E (ti)E (ti)† = I . So, we solve the following system ofequationsW e(ti) = E (ti)†α (ti) ∀i ∈ [m],⇒E (ti)W e(ti) = α (ti) ∀i ∈ [m].Now, let us define β (ti) := E (ti)†α (ti). Then, we get the following linear system ofequations to solve:W[e(t1), e(t2), . . . , e(tm)]=[β (t1), β (t2), . . . , β (tm)].Due to Assumption 1 the query tokens have Kruskal Rank at least n, and m≤ n1 ≤ n,so the query token matrix in the left-hand-side above has a pseudo left inverse.Thus, W has at least one solution with rank at most m. Assuming the solutionis W ∗ ∈ Rd×d , there exists W Kh,W Qh ∈ Rd×dh such that W KhW Q⊤h =W ∗ for theweights of the Attention head. So, we can prove Claim 1 by setting m = n and(t1, t2, . . . tn) = (1,2, . . . ,n).Next, to prove Claim 2, we prove that if m < n1 ≤ n, (t1, t2, . . . , tm) = (T −m+ 1,T −m+ 2, . . . ,T ), then we can not only achieve any desired logits for theexamples (T −m+ 1,T −m+ 2, . . . ,T ), but also keep the {θ (t)}T−mt=1 unchanged.So far, we have found a W ∗ that achieves the set of desired logits. Now, there existsa linear transformation W+ ∈ Rd×d such thatE (t)W+e(t) = 0 for all T −m< t ≤ T, (B.1)[E (t)W+e(t)]i̸=[E (t)W+e(t)]jfor all 1≤ t ≤ T −m, i, j ∈ [n], i ̸= j. (B.2)This is due to Lemma 4.So, we can set W :=W ∗+cW+, and let c→ ∞. This way, the desired logits forT −m< t ≤ T stay the same since cW+ has no effect on them, while for the rest of83logits, cW+ dominates W ∗ and saturates the Softmax operator (Lemma 5). So, wecan prove Claim 2 by finding W Kh,W Qh ∈Rd×dh such that W KhW Q⊤h =W ∗+cW+,and letting c→∞. This construction is feasible because rank(W ∗+cW+)≤m+1≤n1 ≤ dh.B.2 Proof of Proposition 2First, we derive the output tokens E ′(t),e′(t) as a function of input tokens. We omitthe head subscript since we have only one head. Since W K =W Q = 0, we haveα (t) = 0, and consequently θ (t) = 1n 1. Therefore, the attention head only takes theaverage of each token, i.e.,e¯ :=1n(e(t)1 + e(t)2 + · · ·+ e(t)n).Now, since the Attention head also has a skip connection, the output of each tokenbecomes:e′(t)= e(t)+o(t) = e(t)+ e¯, (B.3)e′(t)i := e(t)i +o(t)i = e(t)i + e¯ for all i ∈ [n] (B.4)First, since we have e′(t) = e˜(t) in Assumption 3, this means that the set {e′(t)}Tt=1has a Kruskal Rank at least n, and Assumption 1 holds for T ′. Next, since the E (t)are full-rank, adding the average of tokens to E (t) preserves the rank (Lemma 3), soE ′(t) are also full-rank and Assumption 2 is satisfied for T ′.B.3 Extension of Theorem 1 to Attention with skipconnectionTheorem 1 holds true for when we have Attention with skip connection. Rememberthat skip connection keeps all the computational mechanisms of MHA the sameexcept Eq. (3.6), where the updated equation is:o :=W O⊤[p⊤1 ; p⊤2 ; . . . ; p⊤H]+ e (o ∈ Rde)84The only modification that we have to make, is to modify the labels. Let usconsider the setting of Theorem 1, with dataset T ={(E (t),e(t),y(t))}Tt=1andAttention weights W . Now, if we add the skip connection to the Attention layerwith the same set of weights, the new yˆ(t) will be:yˆ(t) =W D⊤(e(t)+o(t)) = Idout×d(e(t)+o(t)) = y(t)+ Idout×de(t).Therefore, we get labels y(t)+ Idout×de(t), which is the true label shifted by the firstdout dimensions of the query vector. So, we need to define a new datasetT ′ ={(E (t),e(t),y(t)− Idout×de(t))}Tt=1,and find Attention weights W ′ for T ′ according to Theorem 1. Then, we can usethe set of Attention weights W ′ to get the correct label y(t)+ Idout×de(t)− Idout×de(t)for an Attention layer with skip connection and dataset T .B.4 Additional LemmasLemma 2. For any real-valued block matrix X of the formX =[A C0 B]we have rank(X )≥ rank(A)+ rank(B).Lemma 3. Assume vectors e1,e2, . . . ,en ∈ Rd are linearly independent. Then,e′i := ei+1n(e1+ e2+ · · ·+ en),are also linearly independent for all i ∈ [n].Proof. Due to the linear independent assumption, the only solution toa1e1+a2e2+ · · ·+anen = 0 s.t. a1,a2, . . . ,an ∈ Ris a1 = a2 = · · ·= an = 0. Now, let us examine the same system of equations for e′i.85Namely,b1e′1+b2e′2+ · · ·+bne′n = 0 s.t. b1,b2, . . . ,bn ∈ R.We prove zero is the only solution to this system to conclude the proof.b1e′1+b2e′2+ · · ·+bne′n=b1e1+b2e2+ · · ·+bnen+ 1n(b1+b2+ · · ·+bn)︸ ︷︷ ︸b¯(e1+ e2+ · · ·+ en)=(b1+ b¯)e1+(b2+ b¯)e2+ · · ·+(bn+ b¯)en = 0.According to the linear independence assumption, we must havebi+ b¯ = 0 for all i ∈ [n].By summing all the terms up, we conclude that bi = 0 for all i ∈ [n] is the onlysolution and this concludes the proof.Lemma 4. Consider a set of full-rank matrices E (1),E (2), . . . ,E (T ) ∈ Rn×d wheren< d. Moreover, consider a set of vectors e(1),e(2), . . . ,e(T ) ∈ Rd with a KruskalRank at least n. Then, for any m≤ n−1 there exists a rank-one matrix W ∈ Rd×dsuch thatE (t)W e(t) = 0 for all 1≤ t ≤ m, (B.5)[E (t)W e(t)]i̸=[E (t)W e(t)]jfor all m< t ≤ T, i, j ∈ [n], i ̸= j. (B.6)Proof. Let us define a basis for Rd as b1,b2,b3, . . . ,bd ∈ Rd such that bi = e(i) forall i ∈ [m], and bm+1, . . .bm+2, . . . ,bd are an arbitrary extension of e[m] to a basis.Then, a linear transformation W is uniquely determined by defining c1,c2, . . .cd ∈Rd and settingW bi = ci for all i ∈ [d].First, to satisfy Eq. (B.5), we set ci := 0 for all i ∈ [m], which gives W bi =W e(i) =ci = 0 for all i ∈ [m]. Let us decompose every vector e(t) into the mentioned basis,86i.e.,e(t) = a(t)1 b1+a(t)2 b2+ · · ·+a(t)d bd for all m+1≤ t ≤ T.Since {e(t)}Tt=1 have a Kruskal Rank at least n, then at least one of a(t)m+1, . . . ,a(t)m+2,a(t)dmust be nonzero, otherwise, e(t) can be written as a linear combination of the firstm vectors. Moreover, there exists fixed um+1,um+2, . . . ,ud ∈ R such thata(t)m+1um+1+a(t)m+2um+2 · · ·+a(t)d ud ̸= 0 for all m+1≤ i≤ d,for instance, we can choose ui iid from a Standard Gaussian Distribution and satisfythis almost surely. Now, let us definec ∼N (0, I) (B.7)ci := uic for all m+1≤ i≤ d. (B.8)Therefore,W e(t) = a(t)1 W b1+a(t)2 W b2+ · · ·+a(t)d W bd (B.9)= a(t)1 c1+a(t)2 c2+ · · ·+a(t)d cd (B.10)= a(t)m+1cm+1+a(t)m+2cm+2 · · ·+a(t)d cd (B.11)=(a(t)m+1um+1+am+2um+2 · · ·+a(t)d ud)c for all m+1≤ t ≤ T, (B.12)with one nonzero coefficient. Now, since c is standard gaussian, thenW e(t) ∼N(0,(a(t)m+1um+1+am+2um+2 · · ·+a(t)d ud)2I).Moreover, E (t) is full-rank, so none of its rows are equal to each other. As a result,(e(t)i − e(t)j)⊤W e(t)follows a Gaussian distribution for all i, j ∈ [n], i ̸= j. Since {0} has zero Lebesguemeasure, the probability of having two equal elements is zero. So, Eq. B.6 issatisfied almost surely. Therefore, there exists a rank-one matrix W that satisfies Eq.87(B.6) and this completes the proof.Lemma 5. Consider an input matrix E ∈ Rn×d , an input vector e ∈ Rd , and aweight matrix W+ ∈ Rd×d such that[EW+e]i ̸=[EW+e]j for all i, j ∈ [n], i ̸= j. (B.13)For any W ∗ ∈ Rd×d define W c :=W ∗+ cW+ where c ∈ R is a scalar. Then,θ := limc→∞Softmax([EW ce]) ,is only a function of W+.Proof. Let us defineα+ :=[EW+e], and α ∗ := [EW ∗e] .Now, ith element of θ is given by:θ i = limc→∞exp(α ∗i + cα+i)∑ j∈[n] exp(α ∗j + cα+j)= limc→∞exp(α ∗i )exp(α ∗i )+∑ j ̸=i exp(α ∗j + c(α+j −α+i ))= 1[i = argmax{α+1 ,α+2 , . . . ,α+n }],where the last equality is due to α+i being distinct. Therefore, θ i is a one-hot vectordetermined by W+ and this completes the proof.B.5 Upper bound on the rank of matrix Z with strongerassumptionsIn Proposition 1 we provided a lower bound on the rank of Z , which matches theupper bound of Lemma 1. One might wonder whether under stronger assumptions88(and perhaps far from reality), a full-rank Z with rank Hd is achievable or not. Here,we prove that even the "General Position" on all tokens assumption is not enough,and there exists a very simple assumption that limits the rank of Z to H(d−1)+1.Proposition 3. Assume that we have a multi-head attention layer A with H heads,embedding dimensions de = dh = d, and dv = 1. Let T ={(E (t),e(t),y(t))}Tt=1be a training set with context size n < d. If all context vectors satisfy 1⊤e(t)i =Cfor all t ∈ [T ], i ∈ [n] and some universal scalar C ∈ R, then for any set of weights{(W Kh,W Qh)}Hh=1, we haverank(Z)≤ H(d−1)+1.Proposition 3 gives an upper bound on the rank of Z when individual elementsof each context token sum up to a constant C. Note that all the context vectorscan be in general position and still satisfy this condition. Moreover, this upperbound puts no assumption on the query vectors. Hence even with the strongestassumptions, obtaining a full-rank matrix Z is out of reach. However, whetherstronger assumptions can maintain dependence on d instead of n (i.e., give H(d−1)+1 memorization lower bound instead of H(n1−1)+1) is still an open questionfor future work.Proof of Proposition 3. First, note that since 1⊤e(t)i =C, and θ(t)i are convex com-binations we have:1⊤z(t)h = 1⊤E (t)⊤θ (t)h =C1⊤θ (t)h =C for all t ∈ [T ].With some abuse of notation, let us write the feature matrix Z ∈ RT×Hd as follows:Z =[Z1, Z2, . . . , ZH],where each Zh ∈ RT×d is a sub-block of Z corresponding to the h-th head for allh ∈ [H]. Then, since each row of Zh (i.e., z(t)h ) sums up to C, with an argument89similar to the proof of Proposition 1 we get:rank(Z) = rank([Z1, Z2, Z3, . . . , ZH])= rank([Z1, Z2−Z1, Z3−Z1, . . . , ZH −Z1])≤ rank(Z1)+ rank(Z2−Z1)+ · · ·+ rank(ZH −Z1)≤ d+(d−1)+(d−1)+ · · ·+(d−1)= H(d−1)+1.B.5.1 Synthetic experiments under general position assumptionWe also run synthetic experiments under the general position assumption on allinput tokens. We consider a similar setting to Section 3.5.2, but we also generatecontext vectors from a uniform distribution as well, i.e.x(t)1 ,x(t)2 , . . . ,x(t)n ,x(t)iid∼U (0,1)64 for all t ∈ [T ].This way, all tokens (i.e., context and query) inside, and across all examples wouldsatisfy the general position assumption. Preliminary results in Figure B.1 suggestthat the linearity in H still remains, while the linearity in n is weakened. We conjec-ture that a bound with complexity O(Hd) would hold, which could be interestingfor future works. Yet, as we showed in Section 3.5.1, from a practical point of view,the General Position assumption might be too strong to be realistic.B.6 Testing Assumptions on Natural Language TasksWe further test Assumptions 2 and 1 on BERT [27], and GPT2 [75] on the task ofpre-training on Wikipedia [34]. Following the setting introduced in Section 3.5.1,we examine the embedding/first layer of trained/untrained GPT2/BERT.• BERT: We follow a setting similar to the pre-training of BERT. For eachexample, an input of the format “[CLS] Sentence1 [SEP] Sentence2” isformed, where Sentence 1 and Sentence 2 are two pieces of text from the900 10 20 30 40 50 60# Heads (H)0.00.20.40.60.81.0Memorized FractionClassification AccuracyT=130kT=110kT=90kT=70kT=50kT=30k(a)0 10 20 30 40 50 60# Heads (H)10 1310 1110 910 710 510 310 1MSE LossRegression LossT=90kT=70kT=50kT=30kT=10k(b)0 10 20 30 40 50 60Context Size (n)0.30.40.50.60.70.80.91.0Memorized FractionClassification AccuracyT=60kT=50kT=40kT=30kT=20k(c)0 10 20 30 40 50 60Context Size (n)10 1210 1010 810 610 4MSE LossRegression LossT=90kT=70kT=50kT=30kT=10k(d)Figure B.1: Testing memorization under general position assumption for bothcontext and query vectors. While the linearity in H still remains, thelinearity in n no longer strictly holds, suggesting H(d−1)+1 to be abetter alternative under general position assumption.Wikipedia corpus. Here Sentence 2 is the sentence after Sentence 1 withprobability 0.5 and randomly chosen from the corpus with probability 0.5.The “[CLS]” token’s output predicts whether Sentence 2 follows the firstsentence or not. We consider the “[CLS]” token as the query token e(t), andthe whole input as the context E (t) for all t ∈ [T ].• GPT2: This model is trained autoregressively to predict the next token ineach piece of text. We consider last-token prediction from the context. Thatis, we set the query e(t) to be the last token of an example, and the contextE (t) to be the whole input example.The results of testing Assumptions are shown in Table B.2 While linear inde-91Table B.2: Testing assumptions on Wikipedia for BERT and GPT2. L0 refersto the embedding layer, and L1 refers to the output of the first layer.GPT2 BERTModel Random Trained Random TrainedL0 L1 L0 L1 L0 L1 L0 L1General Position × × × × × × × ×Assumption 1 × ✓ × ✓ × × × ×Assumption 2 ✓ ✓ ✓ × ✓ ✓ ✓ ✓pendence assumptions are not well-suited for discrete data, we observe that mostassumptions still hold. In the case of BERT, the n/d ratio is 512/1024, and whilethe Kruskal Rank of query vectors is below 512 (violation of Assumption 1), weobserve that the Kruskal Rank is high enough (291 for the case of trained and349 for the case of random BERT). Hence, the result of Remark 1 still yields anacceptable lower bound.92
Thesis/Dissertation
eng
Electrical and Computer Engineering
Vancouver : University of British Columbia Library
Attribution-NonCommercial-NoDerivatives 4.0 International
Graduate
Liao, Renjie
Thrampoulidis, Christos
Exploring algorithmic reasoning and memorization in transformers : challenges and insights
Text