SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 2 shows the hyperparameters we used in our experiments. We use the same hyperparameters for all models and positional encoding schemes. In our initial experiment, we tried a few more hyperparameters such as  $\mathrm{lr} \in \{0.00001, 0.00003, 0.00005\}$  and WeightDecay  $\in \{0, 0.05, 0.1\}$ , but we did not observe any significant difference in the results. So, we decided to use the same hyperparameters throughout our experiments.
In our experiments, we used single-GPU training setup for the models. Specifically, we ran our experiments on a mix of NVIDIA V100 32G, NVIDIA RTX8000 48G, NVIDIA A100 40G, and NVIDIA A100 80G GPUs. Depending on the GPU type and the positional encoding, each of our training runs took 6 to 15 hours, per each seed, on average to complete. Considering all the datasets, and positional encoding schemes, in addition to the scratchpad experiments, and three seeds, we ran about 870 individual training runs for the results in this paper.
D.4 Reproducibility
In this study, all experiments employed open-source libraries, specifically HuggingFace (Wolf et al., 2020) from which we utilized their implementation as a foundation for the training loop, optimizer, and the Transformer architecture. To ensure reproducibility, we will also release a singularity binary with all dependencies and libraries to enable running our experiments on any machine with NVIDIA GPUs and at any time in the future. Moreover, every reported number in this paper is linked to the source code package that deterministically (up to GPU stochasticity) reproduces the results, which we release publicly on GitHub at https://github.com/McGill-NLP/length-generalization.
E.1 Detailed Model Accuracy
We report the detailed results of our experiments in Figures F.4 to F.6. We refer the readers to Appendix D.1 for the description of each task.
Table 2: Summary of hyperparameters used in the experiments.
Parameter	Value
Optimizer	AdamW
Learning rate	0.00003
Weight Decay	0.05
Batch size	64
Learning Rate Scheduler	Polynomial
Warm Up	6% of training steps
# Train Steps	40K steps
Decoding Method	Greedy (No Sampling)
Dropout (taken from HuggingFace)	0.1
Model dimension (taken from HuggingFace)	768
# Layers (taken from HuggingFace)	12
# Attention Heads (taken from HuggingFace)	12
E.2 Detailed Head Distance
Figure F.7 shows the layer-wise distance of No PE's attention patterns with other positional encoding schemes measured across instances of the SCAN dataset. We refer the readers to Section 5.2 for the details and analysis of these results.
E.3 Detailed Model Accuracy On Various Scratchpad Formats
Figure F.8 shows the generalization of various scratchpad formats for each model aggregated across all datasets. Figures F.9 to F.15 show the generalization of various scratchpad formats for each model on each dataset. We refer the readers to Section 6 for the details and analysis of these results.
F Pretraining at 1.3B Scale
In this section, we elucidate our preliminary assessment of length generalization on a pretrained LLM scaled to 1.3B. Ensuring a fair evaluation, we pretrain models across varied positional encodings, all on identical data with consistent parameters. Given the intrinsic significance of element positions to the semantics in code data, our choice leaned towards code-based pretraining.
F.1 Model Architecture
Our approach remains consistent with the original codebase, merely increasing the model dimensions. Specifically, we adopt a decoder-only Transformer structure encompassing 24 layers. The configuration details are as follows:  $d_{\mathrm{model}} = 1024$ ,  $d_{\mathrm{kv}} = 128$ ,  $d_{\mathrm{ff}} = 16$ , 384, and 32 attention heads. Consequently, this amounts to a Transformer model with 1.3B parameters. We train the models with a context size of 1024 tokens. For the rest of training hyperparameters, we follow Allal et al. (2023). Specifically, we use a global batch size of 256 and AdamW optimizer with  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.95$ ,  $\epsilon = 10^{-8}$ , and weight decay 0.1. Also, the learning rate of  $2 \times 10^{-4}$  with cosine decay and  $2\%$  warm-up is used.
F.2 Dataset Selection
We utilize a subset of 30M documents from the StarCoder dataset (Li et al., 2023), constituting a blend of  $40\%$  Python,  $25\%$  Java,  $25\%$  JavaScript,  $5\%$  GitHub issues, and  $5\%$  GitHub commits. Processed through the StarCoder tokenizer with a vocab size of 49, 152, the dataset results in 30B tokens. We have trained the model for a single epoch. For comprehensive details on the remaining hyperparameters, we direct readers to Allal et al. (2023).
F.3 Generalization Evaluation
Our evaluation aims to assess how the model's performance in predicting a given token varies with changes in context size, especially when the context size exceeds the training context size. Ideally, we expect the perplexity to generally improve as the context size increases since the model has access to more information. In this experiment, we select a sample of 1500 validation documents, each representing a single source code file. For each document, we calculate the perplexity of the last tokens, considering various context sizes ranging from 512 to 2560 (Note the model itself is pretrained with a context size of 1024). This approach allows us to evaluate the model's perplexity under both in-distribution and out-of-distribution conditions. To keep the computation and number of inferences tractable, we measure the perplexity on a maximum of last 200 tokens in the document and ensure exactly the same set of tokens are evaluated per each context size.
To obtain a comprehensive understanding of the model's performance, we also categorize documents into buckets based on their lengths and calculate the average perplexity within each bucket. This categorization enables us to examine the nature of dependencies within each bucket, recognizing that longer documents do not always imply longer dependencies.
We present the results in Figure F.2 and Tables 3 and 4. First, it is observed that all models benefit from a larger context up to the training size of 1024. In the I.I.D. case (context length  $\leq 1024$ ), no significant differences among the models are discernible across all length buckets. Conversely, in the O.O.D. case (context length  $>1024$ ), the model trained with Rotary fails to generalize as its perplexity explodes in all cases, aligning with our findings in Section 4. NoPE and ALiBi effectively generalize to larger context sizes up to 1800. However, as context sizes grow beyond that, ALiBi exhibits relative stability compared to NoPE, though both models show a pattern of increasing perplexity, potentially indicating inefficiencies to capture longer dependencies.
NoPE
Figure F.2: Perplexity of the three pretrained models on examples from various length buckets. Each model is trained with different positional encoding schemes on identical training data and hyperparameters.
Positional Encoding
ALiBi
Table 3: The detailed perplexity (represented along with Figure F.2 for completeness). of the three pretrained models on examples from various length buckets (Part I)
Model	Context Size
Length Bucket [750, 1000]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	2.456	2.314	-	-	-	-	-	-	-	-	-	-
Rotary	2.427	2.310	-	-	-	-	-	-	-	-	-	-
Length Bucket [1000, 1250]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	2.958	2.372	2.282	2.192	-	-	-	-	-	-	-	-
Rotary	2.944	2.392	2.375	2.263	-	-	-	-	-	-	-	-
Length Bucket [1250, 1500]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	3.122	2.951	2.722	2.676	2.618	2.570	-	-	-	-	-	-
Rotary	3.077	2.940	2.816	2.723	2.678	2.670	-	-	-	-	-	-
Length Bucket [1500, 1750]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	2.620	2.541	2.467	2.258	2.187	2.158	2.121	-	-	-	-	-
Rotary	2.616	2.564	2.527	2.305	2.226	2.249	2.307	-	-	-	-	-
Length Bucket [1750, 2000]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	3.002	2.911	2.853	2.704	2.436	2.397	2.399	4.581	-	-	-	-
Rotary	2.966	2.908	2.954	2.774	2.480	2.487	2.595	4.114	-	-	-	-
Length Bucket [2000, 2250]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	4.130	3.998	3.892	3.849	3.751	3.348	3.338	6.325	39.413	-	-	-
Rotary	4.087	3.988	4.049	3.983	3.823	3.479	3.616	5.365	18.731	-	-	-
Length Bucket [2250, 2500]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	2.778	2.642	2.555	2.493	2.450	2.321	2.159	2.859	30.012	146.219	-	-
Rotary	2.764	2.655	2.647	2.570	2.499	2.441	2.361	3.695	14.173	44.089	-	-
Length Bucket [2500, 2750]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	2.805	2.716	2.645	2.614	2.574	2.534	2.367	2.954	25.694	114.870	259.222	-
Rotary	2.801	2.732	2.721	2.726	2.651	2.681	2.575	3.585	12.350	37.746	93.094	-
Length Bucket [2750, 3000]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	2.850	2.776	2.728	2.694	2.666	2.620	2.504	3.839	41.194	151.810	311.138	536.225
Rotary	2.830	2.789	2.825	2.766	2.735	2.742	2.741	3.680	13.814	42.951	102.930	150.322
Length Bucket [3000, 3250]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	2.783	2.697	2.639	2.595	2.555	2.519	2.498	4.239	32.277	153.248	312.356	521.752
Rotary	2.782	2.746	2.750	2.688	2.665	2.692	2.763	3.847	13.542	41.208	97.685	143.644
Length Bucket [3250, 3500]
NoPE	512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
ALiBi	3.260	3.140	3.061	3.009	2.878	2.828	2.866	6.048	36.581	143.613	296.817	479.934
Rotary	3.206	3.116	3.127	3.073	2.922	2.940	3.104	5.077	16.838	47.832	106.607	151.288
Table 4: The detailed perplexity (represented along with Figure F.2 for completeness). of the three pretrained models on examples from various length buckets (Part II)
Model	Context Size
Length Bucket [3500, 3750]
512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
NoPE	3.070	3.003	2.969	2.942	2.894	2.826	2.816	3.808	30.840	122.338	278.935	473.120
ALiBi	3.054	3.011	3.069	3.006	2.954	2.971	3.107	4.845	15.691	46.406	106.966	156.849
Rotary	3.020	2.958	2.923	2.892	2.844	3.719	169.445	604.442	771.338	766.585	833.169	888.627
Length Bucket [3750, 4000]
512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
NoPE	2.703	2.593	2.564	2.517	2.466	2.440	2.427	3.447	35.851	135.291	281.145	438.679
ALiBi	2.657	2.570	2.614	2.553	2.515	2.541	2.639	4.049	14.112	40.070	94.953	138.232
Rotary	2.645	2.536	2.506	2.461	2.416	2.981	147.012	508.386	671.961	707.306	861.924	873.948
Length Bucket [4000, 4250]
512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
NoPE	3.096	3.021	2.947	2.905	2.862	2.843	2.809	3.459	27.171	109.982	244.809	414.226
ALiBi	3.086	3.040	3.064	3.028	2.931	2.983	3.058	4.316	13.509	37.900	88.864	129.419
Rotary	3.054	2.973	2.906	2.868	2.831	3.626	158.974	723.175	1114.522	989.245	1057.793	960.442
Length Bucket [4250, 4500]
512	640	760	878	1024	1200	1400	1600	1800	2048	2304	2560
NoPE	3.056	2.992	2.921	2.869	2.815	2.786	2.872	8.839	42.054	133.110	255.307	418.937
ALiBi	3.004	3.008	3.022	2.975	2.913	2.975	3.097	4.913	15.522	43.614	98.910	142.342
Rotary	2.958	2.899	2.838	2.775	2.731	3.603	174.717	528.865	694.649	741.368	806.988	857.006
Figure F.3: The learnt bias value  $b_{\mathrm{bucket}(.)}$  (taken from the pretrained T5-3b model) visualized for relative distance bucket bucket(.) ∈ [0, 31]) and 12 heads. Bucket 31 is exclusively used for all distances larger than 128 tokens. These results demonstrate that while T5 has heads that mask distant tokens, other heads attend to distant tokens when required.
Problem Length
NoPE
T5's Relative PF
AIjBi
Absolute Position Embedding
Figure F.4: Generalization behavior of positional encoding schemes on Primitive tasks.
Positional Encoding
Figure F.5: Generalization behavior of positional encoding schemes on Mathematical & Reasoning tasks.
Positional Encoding
NoPE
AIiBi
Absolute Position Embedding
Figure F.6: Generalization behavior of positional encoding schemes on Classic Length Generalization tasks.
Example Length
Positional Encoding  $(M)$
NoPF
Figure F.7: Layer-wise distance of No PE's attention patterns with other positional encoding schemes measured across instances of SCAN dataset.
Figure F.8: The optimal scratchpad format is aggregated across all datasets per each model. The optimal scratchpad format is different for each model.
Figure F.9: Generalization of various scratchpad formats for each model on the Addition task.
Figure F.10: Generalization of various scratchpad formats for each model on the Summation task.
Figure F.11: Generalization of various scratchpad formats for each model on the Parity task.
Figure F.12: Generalization of various scratchpad formats for each model on the Sorting task (Single Digit).
Figure F.13: Generalization of various scratchpad formats for each model on the Sorting task (Multi Digit).
Figure F.14: Generalization of various scratchpad formats for each model on the Polynomial Evaluation task.
Figure F.15: Generalization of various scratchpad formats for each model on the LEGO task.