Physics of Language Models: Part 3.1, Knowledge Storage and Extraction — part 5/6

Cleaned-text code points [63018, 78898). Bibliography omitted here; source.md includes it.

## Appendix B Details on Model Architecture

The GPT2-small architecture [ 31 ] has 12 layers, 12 heads, and 768 = 12 × 64 768=12\times 64 hidden dimensions (124M). Recent research [ 16 , 34 , 8 ] has shown that transformers can achieve a significant performance improvement by utilizing attentions based on the relative positional differences of tokens. (This is more systematically studied in [ 1 ] .) Consequently, in this paper, we replace the positional embedding with a rotary embedding, following the standard GPT-NeoX implementation [ 8 ] available on Huggingface (with the default frequency base of 10,000 and rotary dimension set to a 1/4 of the embedding dimension). We continue to refer to this as GPT2 for brevity.

In our bioS experiments, we employ the above architecture. For the bioR experiments, we opt for a larger GPT2 model with 12 layers, 20 attention heads each 64-dimensional (302M), tailored to its increased difficulty. Only when presenting our negative result in Figure 2 , we also tried a 12-layer, 32-head (each 64-dimensional) GPT2 model (682M).

In our knowledge extraction experiments (e.g., Figure 3 ), we also used a downsized Llama architecture with the same number of layers, heads, and hidden dimensions as the GPT2 architecture, yielding similar results. For the probing experiments, we exclusively used the GPT2 architecture for conciseness.

Additionally, we evaluate the BERT model [ 20 ] . BERT is similar to GPT2 but features a complete attention matrix, enabling every token to attend to all others. For a strong side-by-side comparison, we modify our GPT2 architecture to swap its triangular attention matrix for a full matrix, while keeping the GPT2 tokenizer and rotary embedding (removing positional embedding). We label this revised model GBERT. A primary distinction is that GBERT adopts pre-layernorm (inherited from the base GPT2 architecture), whereas BERT utilizes post-layernorm.

Throughout pretraining, mixed training, and QA finetuning, we maintain a context window length of 512.

## Appendix C Details on Pretrain and Mixed Training

During BIO pretraining, we randomly sample biographical entries of individuals and concatenate them to form sequences of 512 tokens, using a standard <EOS> token to separate individual entries.

In mixed training, we pre-train the model with BIO data from all individuals and QA data from half of them. Specifically, each training sequence of 512 tokens is either sourced entirely from the BIO entries (as previously mentioned) or entirely from the QA entries (again, from randomly sampled individuals and concatenated). We define a parameter 𝖰𝖠 r {\mathsf{QA}_{r}} to dictate the frequency of using QA entries. Predominantly in this paper, we set 𝖰𝖠 r = 0.8 {\mathsf{QA}_{r}}=0.8 , which implies a 2 : 8 2:8 ratio between BIO and QA entries in terms of the number of pre-trained tokens. We subsequently assess the model’s generation accuracy using QA data from the other half of the individuals. Refer to Figure 10 for an analysis of how the parameter 𝖰𝖠 r {\mathsf{QA}_{r}} impacts mix-training performance.

For both BIO pretraining and mixed training, we employed a conventional set of optimization parameters: the AdamW optimizer with a weight decay of 0.1, ε = 10 − 6 \varepsilon=10^{-6} , an initial learning rate of 0.001, a 1000-step linear warmup, and cosine learning rate decay (from 0.001 decreasing to 0.0001). We used a batch size of 96.

There were a total of 80,000 training steps for bioS (using the 12-layer, 12-head GPT2/Llama architecture) and 150,000 training steps for bioR (using the larger 12-layer, 20-head GPT2/Llama). Only when using the 12-layer, 32-head GPT2 to present our negative result in Figure 2 , we used 200,000 training steps.

###### Remark C.1 .

Our training time is long enough to ensure next-token prediction accuracy well above 99 % 99\% for both BIO pretraining and mixed training, when focusing on tokens describing six attributes per individual. These numbers are not included in this paper’s figures.

## Appendix D Details on QA Finetune

In our QA finetuning tasks, we first use a BIO pretrained model checkpoint and then apply either full finetuning or LoRA finetuning.

For full finetuning, we employ the AdamW optimizer with ε = 10 − 6 \varepsilon=10^{-6} . We use weight decays of 0.01 and 0.001, and initial learning rates of 0.001 , 0.0003 0.001,0.0003 , and 0.0001 0.0001 . There is no warmup, and we implement cosine learning rate scheduling (reducing to 10 % 10\% of the initial learning rate), a batch size of 48, and a total of 50,000 training steps. Given that we are presenting a negative result for full finetuning (as seen in Figure 2 ), we display the best QA test accuracy among all the lr/wd parameter combinations.

For LoRA finetuning, we maintain the aforementioned AdamW configuration but set a consistent weight decay of 0.01 and an initial learning rate of 0.0003 0.0003 for all tasks.

The results in Figure 11 suggest that for the purpose of QA finetuning, LoRA is generally a better option compared to full finetuning. While a large rank- r r update on the query/value matrices isn’t essential, it appears beneficial to have a significant rank- r ′ r^{\prime} update on the embedding layer to address the distribution shift from the BIO data to the QA data.

For this reason, in all subsequent experiments in this paper (notably Figure 3 and 12 ), when conducting QA finetuning, we use r ′ = 128 r^{\prime}=128 and either r = 8 r=8 or r = 16 r=16 , presenting the best accuracy from the two runs.

## Appendix E Details on P-probing

In our P-probing experiments, we freeze the BIO pretrained GPT2 model and append a limited set of trainable parameters. Using the GPT2-small as an example, we introduce: • a trainable rank-2 update for the embedding layer, having dimensions of 50256 × 2 50256\times 2 and 2 × 768 2\times 768 ,

• for each prediction task that is an M M -class classification problem, a trainable linear layer with dimensions of 768 × M 768\times M ,

• preceding the linear layer, a layer normalization layer furnished with trainable affine parameters.

In the context of P-probing, recall that we considered six classification sub-tasks (from 6 special locations) for every attribute prediction task. Specifically, for the birthdate attribute, we solely address its first-token prediction task, which is equivalent to predicting the individual’s birth month. 25 25 25 This is because a birthdate encompasses 200 × 12 × 28 200\times 12\times 28 potential choices, surpassing N / 2 N/2 , the number of training individuals. For the remaining five attributes, both the first-token and whole-attribute prediction tasks are examined. In sum, this results in 11 prediction tasks, each comprising 6 sub-tasks. For every one of these 11 tasks, we incorporate a distinct set of trainable parameters.

For optimization, the AdamW optimizer is employed with ε = 10 − 6 \varepsilon=10^{-6} , weight decay of 0.3, an initial learning rate of 0.001 0.001 , no warmup, and a linear learning rate decay (down to 0 0 in the end). We set the batch size of 50 and trained for 30,000 steps. During this P-probing training phase, we have turned on the dropout on the (frozen) pretrained GPT2 model to prevent overfitting.

We perform experiments on both bioS and bioR data families (refer to Figure 13 and Figure 14 ), evaluating the P-probing accuracy of first-token and whole-attribute predictions. These figures also compare rank-2 and rank-4 updates on the embedding layer, demonstrating that a large modification to this layer is not crucial for P-probing attribute values.

### E.1 Details on Closer P-Probing

In Figure 6 and Section 5.1.1 , we examined the P-probing results using a Venn diagram based on the bioS couple dataset from Appendix A.1.1 , where each individual has 2 biographical entries. Figure 15 supplements this with results for individuals with 1 or 5 entries.

## Appendix F Details on Q-probing

Recall that in Q-probing, we freeze the pretrained GPT2 model and append a small set of trainable parameters on top for probing purposes. Using GPT2 small as an example, we add: • a trainable rank- r r update on the embedding layer with dimensions of 50256 × r 50256\times r and r × 768 r\times 768 ,

• a trainable linear layer with dimensions of 768 × M 768\times M for each prediction task that is an M M -class classification problem,

• a batch normalization layer before the linear layer, with trainable affine parameters.

We consider an input sentence that only contains a person’s full name, preceded by a starting token and followed by an ending token. After applying all 12 layers of GPT2, we extract the hidden states from the last layer at the ending token. For instance, in the GPT2-small model, this is a 768-dimensional vector. We then apply a linear classifier on top to predict the person’s attributes. Similar to P-probing, we adopt a separate set of trainable parameters for each of the 11 classification tasks.

We employ the AdamW optimizer with ε = 10 − 6 \varepsilon=10^{-6} , a weight decay of 0.3, an initial learning rate of 0.001 0.001 , no warmup, and a linear learning rate decay schedule (reducing to 0 0 by the end). The batch size is set to 200, and we run a total of 30,000 training steps. During training, we allow the frozen GPT2 model to use dropout.

Experiments are conducted on both the bioS and the bioR data families, as shown in Figure 16 , for first-token prediction and whole-attribute prediction. We compare rank-16 versus rank-64 updates on the embedding layer for the bioS data (or rank-32 versus rank-128 updates for the bioR data). This demonstrates that for Q-probing, a larger modification to the embedding layer is not necessary to probe the desired attribute values.

## Appendix G Details on Celebrity Augementation

Recall that in the celebrity knowledge augmentation, we introduced an additional set of N = 100,000 N=100,000 individuals and designated them as the celebrity group, 𝒫 𝖼𝖾𝗅 {\mathcal{P}_{\mathsf{cel}}} . In contrast, the original N N individuals represent the minority group, 𝒫 𝗆𝗂𝗇 {\mathcal{P}_{\mathsf{min}}} . There is no overlap between these two sets of individuals; specifically, they have distinct full names.

In the main body of this paper (specifically in Figure 8 ), we considered two choices: • The minority uses bioS single+permute1 , and the celebrity uses bioS multi5+permute . We denote this combination as bioS single+permute1+CEL and compare it to bioS single+permute1 .

• The minority uses bioR single , and the celebrity uses bioR multi5 . We denote this combination as bioR single+CEL and compare it to bioR single . (We also compare the latter to bioR single+wiki . By this, we mean that during BIO pretraining, half of the training sentences come from the WikiBook dataset, while the other half come from the bioR single data.) 26 26 26 Recall that BERT and RoBERTa were trained on a combination of BookCorpus [ 41 ] and English Wikipedia, which totals 16GB of uncompressed text [ 20 , 25 ] . We use this same 16GB WikiBook dataset.

Note that in both cases, each individual in the minority group has only one biographical entry, while each individual in the celebrity group has five biographical entries. Thus, during BIO pretraining, the BIO data on 𝒫 𝖼𝖾𝗅 {\mathcal{P}_{\mathsf{cel}}} appear with a 1 / 6 1/6 chance.

In this appendix, we explore a broader set of augmentation options. • The minority uses bioS single and the celebrity uses bioS multi5+permute , denoted as bioS single+CEL . We compare this to bioS single . In this scenario, the celebrity and minority groups have biographical entries in different formats : the entries of the celebrity group are randomly shuffled , while those of the minority group follow a fixed order (see ( 2.1 ) ). The QA test accuracy on the minority group increases with the addition of the celebrity group, but not to the same extent as in the bioS single+permute1+CEL case.

• The minority uses bioS single+permute1+fullname and the celebrity uses bioS multi5+permute , denoted as bioS single+permute1+fullname+CEL . We compare this to bioS single+permute1+fullname . In this scenario, the celebrity and minority groups have their biographical entries in different formats : the minority group uses the fullname augmentation, repeating the individual’s full name in each sentence, while the celebrity group only mentions the fullname once. The QA test accuracy on the minority group increases with the assistance of the celebrity group, but not as much as in the bioS single+permute1+CEL case.

• The minority uses bioR single+fullname and the celebrity uses bioR multi5+fullname , denoted as bioR single+fullname+CEL . We compare this to bioR single+fullname . In this case, the celebrity and minority groups have their biographical entries in the same format , leading to a significant increase in QA test accuracy to 82.2 % 82.2\% . (We also compare this to bioR single+fullname+wiki , where during BIO pretraining, half of the training sentences come from the WikiBook dataset, and the other half from the bioR single+fullname data. C.f. Remark 6.1 )

The transformer model is pretrained on the combined set of biographies 𝒫 𝖼𝖾𝗅 ∪ 𝒫 𝗆𝗂𝗇 {\mathcal{P}_{\mathsf{cel}}}\cup{\mathcal{P}_{\mathsf{min}}} and then finetuned using QAs from the celebrity group 𝒫 𝖼𝖾𝗅 {\mathcal{P}_{\mathsf{cel}}} . We evaluate the model’s QA generation accuracy on the 𝒫 𝗆𝗂𝗇 {\mathcal{P}_{\mathsf{min}}} group. 27 27 27 We also considered other fine-tuning variations, such as QA finetuning with half of 𝒫 𝗆𝗂𝗇 {\mathcal{P}_{\mathsf{min}}} as training and half as testing, but found negligible differences. Our findings are reported in Figure 17 .

P-probing and Q-probing. We incorporate P-probing and Q-probing results for our celebrity case. The inclusion of celebrity data enhances the model’s structural knowledge storage, even for minority groups . Figure 18 demonstrates that knowledge about minority groups is often stored in earlier tokens. This confirms that for minority groups , individual full names can more directly encode the six target attributes, due to the introduction of celebrity data. This accounts for the high knowledge-extraction QA accuracies.

## Appendix H Details on BERT Experiment

Recall that GBERT is a bi-directional variant of GPT2, using the same tokenizer, as detailed in Appendix B . It is similar to BERT, but its architecture closely resembles GPT2 for a more direct comparison. We use GBERT for the following tasks: (1) BIO pretrain, (2) BIO+QA mixed training, (3) QA finetune from BIO pretrain, and (4) Q-probing from BIO pretrain. Since we only apply GBERT to the bioS data family to demonstrate a negative result, we utilize the same architecture size as GPT2-small.

For BIO pretrain and BIO+QA mixed training, we use the AdamW optimizer with weight decay 0.1, ε = 10 − 6 \varepsilon=10^{-6} , an initial learning rate of 0.0003, a 1000-step linear warmup, and cosine learning rate decay (from 0.0003 to 0.00003). We use a batch size of 96 for 150000 training steps on the bioS dataset. This is twice the training time compared to the 80000 steps used for GPT2 small on the same dataset, as we are presenting a negative result on GBERT. For BIO+QA mixed training, we tested both 𝖰𝖠 r = 0.2 {\mathsf{QA}_{r}}=0.2 and 𝖰𝖠 r = 0.8 {\mathsf{QA}_{r}}=0.8 and report the best test accuracy.

For QA finetune, we tested four LoRA variants and report their best accuracy. 28 28 28 Specifically, we tested rank-8 or rank-32 update on the query/value matrices, and rank-128 update or full fine-tuning on the embedding layer. We use the AdamW optimizer with weight decay 0.01 and an initial learning rate of 0.0003 0.0003 for all tasks, with linear learning rate decay (down to 0). We use a batch size of 48 for 50000 training steps.