Knowledge Entropy Decay during Language Model Pretraining Hinders New Knowledge Acquisition

Dear Reviewer yokT,

Thank you for your valuable time in reviewing our paper. We appreciate recognizing the novelty of our work and experimental results.

The usage of the term "entropy" in this case can be confusing in my opinion. Usually, this term implies some form of "surprise" with respect to information. However, in this case, it seems to refer to the idea of source broadness.

The reviewer raises a question about whether the observed behavior regarding the knowledge entropy holds across different models, architectures, or training time design choices. To address this, we further test the behavior of knowledge entropy using the Pythia model (1.4b) with Dolma, we observed that the behavior remains consistent. We add the result of Pythia knowledge entropy and update the paper where the figure is in Figure 6 of the revised version of the paper. We kindly ask the reviewer to check the figure of the trend from the updated paper as we cannot attach figures in the rebuttal. Also in the original and the revised paper, to investigate whether the findings hold for various cases, we conducted experiments where we changed the activation function from SwiGLU to ReLU (Figure 9) and adjusted model sizes (1b, 7b) in Figure 2.

Due to computational constraints and the lack of models that provide intermediate checkpoints, we were unable to explore all possible training time design choices. However, based on the results from the various combinations we tested, we expect that the general trend would persist across different training scenarios.

Further discussion is required regarding the correspondence of knowledge entropy and forgetting. Rather than being a direct consequence of knowledge entropy, can it be the case that the reason behind sharper declines in knowledge retention is due to the model's richer knowledge representation between concepts at later stages of the training and thus being more prone to knowledge loss when small changes are made? More specifically, forgetting at a higher rate at later stages can be expected simply due to having more knowledge encoded in comparison to having very little knowledge in earlier stages.

The reviewer suggests that the decline in knowledge retention may be due to the model’s richer knowledge representation at later stages of training since the later models are likely to contain more knowledge compared to initial models. We agree that models in later stages of training would typically encode more knowledge than earlier stages. However, we believe the key factor is not just the amount of knowledge, but how it is stored and utilized.

Our resuscitation experiment highlights this distinction. By altering the knowledge entropy while keeping the same parameter state (the final checkpoint), we observed significant differences in knowledge retention. This suggests that how knowledge is stored and accessed plays a crucial role. Specifically, when the model relies on a small set of certain memory vectors, the rate of forgetting is higher.

Minor grammatical errors that can be fixed in the camera-ready version.

Thank you for pointing this out. We will address the minor grammatical errors in the camera-ready version.

I believe that Algorithm 1 can be rewritten to better relay the information contained. Namely, by giving a better explanation of each operation (especially with respect to t).

Thank you for the feedback. We will revise Algorithm 1 to make it clearer and more explainable, particularly with respect to t, as the reviewer suggested.

Given that consistent behavior is observed across all layers, do you think we can use only a subset of layers to determine the model's proneness to learning new knowledge? For example, only focusing on layers that mainly encode semantic knowledge as shown in the literature.

The reviewer inquires if the behavior of a subset of layers can determine the model’s ability to learn and retain knowledge, especially in terms of the critical layers that have been proposed as encoding semantic knowledge.

According to Meng et al. (2022)[1], feed-forward layers in the middle range of a model play a crucial role in processing the last token of the subject name for factual mediation. Meng et al. (2023)[2] and Meng et al. (2022) identified this range as layers 4 to 9 in GPT-J (6B), which has 28 layers, and up to layer 15 in GPT-2 XL (1.5B), which has 48 layers.

As shown in Figure 8 (Appendix A.2 in the revised version of paper), the layers identified as critical for processing factual knowledge (3rd to 8th layers) exhibit the highest values of Knowledge Entropy and experience less reduction in knowledge entropy during pretraining. The overall reduction in knowledge entropy was driven by the other layers. Based on this observation, it seems that relying on only a subset of layers may not be representative enough to assess the model's overall ability.

That part as it stands leaves this reader wanting more information. For instance, what is the target span for the probes? Also is \theta_{pt} unique? Is there only one such check point and what about \theta_{cl}? Is \theta_{cl} when the continual learning stops, if it does, or what? Is \theta_{cl} then the model after pre training? There is more material in the appendix on this but it would help the reader's understanding of this important concept of knowledge acquisition and knowledge loss to move some of that discussion into the body of the paper.

The reviewer raises a question about lack of details in knowledge acquisition measurement.
To evaluate the model's ability to acquire new knowledge while retaining pre-existing capabilities, we further train the intermediate checkpoints of OLMo within a continual knowledge learning setup. Here, $\theta_{pt}$​ represents the model at a specific stage of pretraining, defined by the number of steps completed during pretraining, while $\theta_{CL}$​ denotes the model after undergoing continual knowledge learning. The knowledge acquisition and retention capabilities of each model are assessed by comparing the performance differences between $\theta_{CL}$ and $\theta_{pt}$. We will include additional details with examples in the next version to clear our explanation of the knowledge acquisition measurement method. We will also ensure the main text contains enough information for clarity.

Some minor stylistic corrections that will render the paper more readable:

Thank you sincerely for providing such detailed and thoughtful feedback, even on minor aspects. We will look over the minor stylistic errors for the camera-ready version.

References

[1] Transformer feed-forward layers are key-value memories. Geva et al.,

[2] OLMo: Accelerating the Science of Language Models Groeneveld et al.,

[3] How Do Large Language Models Acquire Factual Knowledge During Pretraining Chang et al.,

While it's pretty clear intuitively what attention entropy or next token prediction entropy are, it's less clear at least to this reviewer what the intuition is behind defining entropy in terms of the coefficients of the outputs of a given FFN layer in the model. How does the knowledge entropy so defined link up with the suggestion that "models in the later stage of pretraining tends to engage with a narrower range of memories, relying more heavily on specific memory vectors rather than accessing and integrating knowledge from a broader range of memories"?

The reviewer raises a question about the intuition behind defining entropy in terms of coefficients and how this definition connects to the observed behavior of the model. Building on top of previous research(Geva et al. (2021) [1]) viewing FFN layer as key-value memory, we wondered how the behavior of selecting values evolve as the model attains more parameteric knowledge during pretraining. Specifically, if more knowledge is stored in the last matrix of FFN layer as pretraining progresses, the model trained with more tokens will select relevant memories with more certainty. Thus we proposed knowledge entropy to measure this certainty, spiky distribution of “keys”.

In the paper, as shown in Figure 1, we referred the memory vectors to the last projection matrix of FFN layer and the memory coefficients are input to the last projection. Knowledge entropy measures the sparsity of the coefficients in an FFN layer; while higher knowlege entropy implies uniform (less certain) distribution of the coefficients, integrating broader range of memory vectors, lower knowlege entropy captures spiky(more certain) distribution, focusing on certain range of memory vectors.

The mechanism behind correlation of Knowledge Entropy and capability to acquire and retain knowledge is that the gradient of the loss is propagated to the memory vectors in the FFN layer in proportion to these coefficients. As the coefficients become more concentrated on specific positions, they may repeatedly overwrite the highly activated regions of the parameters rather than distributing across all parameters, thereby limiting the model’s overall knowledge capacity. Our findings reveal a consistent decrease in Knowledge Entropy as pretraining progresses, indicating that the model increasingly relies on certain memory vectors. When these late-stage model checkpoints are trained continually, the overall performance of knowledge acquisition and retention can be constrained due to the limited capacity. We included the intuition behind the definition of Knowledge Entropy and its impact on training in the Appendix A.1 in the revised version of the paper.

Another element that needs a bit more work is the discussion of knowledge acquisition and forgetting. This relies on work from Chang et al. but to be more self-contained the authors could explain in a bit more detail what the probes are doing (and why, for instance, 15 different probes?), or rather summarize Chang et al's contributions more clearly with respect to the contributions of the present work.

The reviewer raises a question about the distinct contribution of our work from Chang et al [2]. Chang et al established a framework to measure LLM’s knowledge acquisition, and as we are measuring the model’s knowledge acquisition, we followed their dataset and evaluation metric. But please note that there is a difference in the training setups due to distinct motivations: while Chang et al analyzes how LLMs acquire knowledge during pretraining, we are rather interested in why LLMs lose this ability during pretraining.

Here, we will briefly explain the knowledge injection and knowledge acquisition measurement framework established by Chang et al. The Fictional Knowledge dataset consists of two components: a fictional corpus, which contains new knowledge presented in a Wikipedia-style format, and corresponding probes that are cloze-task-style sentences to query the information within the corpus. The final span of each probe, referred to as the target span, is used to evaluate the model’s prediction probability, which serves as a measure of knowledge acquisition performance.

In Chang et al., the probes are divided into three levels of difficulty, with five sentences created for each level. This results in 15 probes per corpus. The difficulty levels are as follows: 1) Memorization probes directly ask about sentences explicitly present in the fictional corpus. 2) Semantic generalization probes are paraphrased versions of the memorization probes to test the model’s understanding of meaning beyond surface forms. 3) Compositional generalization probes are designed to assess whether the model can integrate multiple pieces of knowledge from the fictional corpus. We have revised the detailed explanation of the evaluation dataset in Appendix B.1.

Dear reviewer EdjN,

Thank you sincerely for taking the time to share your valuable insights. We appreciate your recognition of the importance of our work and its potential to contribute significantly to the understanding of model training. Your acknowledgment of the solid empirical results and the potential for substantial contributions to the community is truly encouraging. The authors will incorporate your comments to improve the quality of the paper.

The talk of broad and particular memory vectors in the introduction is a bit unclear.

The reviewer highlighted the ambiguity in the concept of broad and particular memory vectors mentioned in the introduction. Our intention was to introduce the idea that a model's processing can be understood as integrating parametric memories, which we later define as memory vectors(the row vectors of the final matrix in feed-forward layers), building on prior literature that interprets feed-forward layers as key-value memories ([1] Geva et al., 2021). By broad and particular memory vectors, we were referring to the model's behavior of whether the memory vectors are combined in uniform distribution (most memory vectors are used with low coefficient) or skewed distribution (only part of the memory vectors are used with high coefficient), respectively. We will clarify the definition of these terms in the final version to ensure better understanding.

In calculating entropy, the authors assume the feed forward network has the form FFN(x) = f(x · K^T ) · V but this leaves out important bias terms that might perhaps affect the calculations. Can the authors say more about this?

The reviewer raises a question regarding the exclusion of the bias term in our entropy calculation; how would the calculation change with inclusion of bias term. We would like to clarify that the model we used in our experiments, OLMo, does not include a bias term. As noted in the OLMo paper[2]: “No biases. Following LLaMA, PaLM, and others, we exclude all bias terms from our architecture in order to improve training stability.”

We believe that even with a bias term, the intuition and interpretation regarding knowledge entropy remain unchanged. Specifically, the output of $f(\mathbf{x} \cdot K^T)$ represents the weights for the linear combination of the corresponding vectors in the $V$ matrix and knowledge entropy measures the distribution of these weights. With a bias term, a memory vector always having a coefficient as 1 is added, which should not affect the distribution of the other coefficients. Thereby, calculation of knowledge entropy does not take into account the bias vector.

Is the resuscitation method intended solely as a test for your main finding, or do you also envision a path in which this could practically be useful for increasing plasticity of models?

The reviewer asks whether the resuscitation method can be viewed as a practical method to increase the plasticity of models. While we observe that our resuscitation method reduces the forgetting rate and improves knowledge acquisition thereby supporting our main findings of the correlation between knowledge entropy and knowledge acquisition/forgetting rates, the primary motivation of the experiment was to test our hypothesis. As such, we did not explore alternative methods for manipulating model parameters and instead focused on a basic approach (Algorithm 1). Thereby, as noted in the paper, we observe that the end performance in knowledge acquisition and retention tends to be higher when selecting a middle-stage checkpoint with similar knowledge entropy, rather than achieving the same entropy through resuscitation. However, we view this approach as an initial attempt towards increasing model plasticity. We believe that refining the parameter manipulation approach beyond the simple technique used in Algorithm 1 could yield better overall performance, especially with larger q values, and think that this approach can be practically used to address major challenges in training language models with new knowledge in the future.

reference

[1] How Do Large Language Models Acquire Factual Knowledge During Pretraining? Chang et al.

[2] Amuro & Char: Analyzing the Relationship between Pre-Training and Fine-Tuning of Large Language Models, Sun and Dredze.

It would also be useful if the authors provide some discussion regarding mechanisms behind their main finding (knowledge entropy correlates with decreased ability to acquire new knowledge). It does not need to be theorem based, which is a difficult endeavor. Empirical findings would help, to delve at least a bit deeper into this phenomenon and why it might occur.

Why do you think that knowledge entropy correlates with decreased ability to acquire new knowledge?

The reviewer raises a question about the mechanisms behind the relationship between decreasing knowledge entropy and the ability to acquire and retain knowledge during pretraining. Our analysis is based on the intuition that the coefficients in the linear combination of memory vectors determine how the corresponding memory vectors are updated.

As defined in Equation 1, the FFN output ($FFN(\mathbf{x}) = f(\mathbf{x} \cdot K^T) \cdot V$) is a linear combination of memory vectors $\mathbf{v}_{i=1,\cdots,m} \in \mathbb{R}^{d}$

(the row vectors of $V \in \mathbb{R}^{m \times d}$ ), where the coefficients $c_i$ are given by $f(\mathbf{x} \cdot K^T)$. In other words, $FFN(x)=c_1\mathbf{v_1}+c_2\mathbf{v_2}+\cdots+c_m\mathbf{v_m}$. Within a given layer, since the operations beyond the FFN layer and the input to the FFN layer remain consistent across all memory vectors, the coefficients act as scaling factors for the gradient by the chain rule. During training, the gradient $\frac{ \partial L}{\partial \mathbf{v_{i,j}}}$ for $i = 1, 2, \ldots, m$ and $j = 1, 2, \ldots, d$ can be decomposed as: $\frac{ \partial L}{\partial \mathbf{v_{i,j}}} = \frac{ \partial L}{\partial FFN(x)}\cdot \frac{ \partial FFN(x)}{\partial v_{i,j}}$

Here, $\frac{ \partial L}{\partial FFN(x)}$ is the same for all $\mathbf{v_i}$, meaning that the relative magnitude of the gradient depends on $\frac{ \partial FFN(x)}{\partial v_{i,j}} = c_i$. Thus, larger coefficients result in proportionally larger gradients being applied to the corresponding memory vectors, amplifying their updates during backpropagation.

As pretraining progresses, a spikier coefficient distribution—captured by decreasing knowledge entropy—implies that gradient updates become increasingly concentrated on specific positions where the average coefficients are larger. This centralization can affect the model's ability to evenly utilize its memory capacity, impacting knowledge acquisition and retention.

We included the intuition in our paper in Appendix A.1. We sincerely appreciate your valuable feedback.

Dear reviewer zgzo,

The authors thank you sincerely for taking your valuable time to provide constructive feedback on this paper. We appreciate your acknowledging the novelty of our defined knowledge entropy, which combines the concept of entropy with key-value memory and the validity of the experiment to deliver our main claim. We present additional experiments and clarification regarding the reviewer’s concerns and questions below.

This work only tested on OLMo models (1B and 7B), as acknowledged by authors. The authors state their reasons for this as due to the fact that OLMo provide intermediate checkpoints. However the Pythia models, too, offer intermediate checkpoints. If time permits, the authors should endeavor to check their results on these models too, in order to ascertain the generality of their claims and whether it holds beyond a single model family.

The reviewer raises a question about whether our findings are applicable to models beyond OLMo, such as Pythia. While we acknowledged that Pythia also offers intermediate checkpoints, we encountered challenges when analyzing its training dynamics. Specifically, Pythia models are trained on a much smaller dataset—approximately 10% of the OLMo training dataset. However, we observe that OLMo models show more stable training dynamics after about 15% of the training dataset has been processed, which is after the last checkpoint of Pythia.

The early stages of training are influenced by a larger proportion of randomly initialized parameters, exhibit different behavior compared to models that have gone through more training steps and reached a stable state. Thereby, although Pythia provides the intermediate checkpoint, they all tend to be before reaching a stable state when analyzed with OLMo models, making Pythia not ideal for analyzing the training dynamics of language models. We believe this is one reason why previous work, such as [1] Chang et al. (2024) and [2] Sun and Dredz(2024), also focused on OLMo models for intermediate checkpoint analysis.

As we cannot include figures in the rebuttal, we add the figure comparing the training dynamics of Pythia and OLMo in the revised version of our paper. Please refer to Figure 18 in the revised version of the paper for the result. Our analysis shows that while the overall training dynamics of Pythia are similar to those of OLMo, Pythia's early-stage training ends prematurely, before significant knowledge acquisition occurs. Additionally, we have included the performance of OLMo at 18k steps in our revised results to show the training dynamics before it stabilizes, which we estimate occurs at around 118k steps.

Dear reviewer zgzo,

We appreciate your time and effort in reviewing the paper.

As we near the end of the discussion period, we would like to remind you of the upcoming deadline kindly. We are happy to discuss any aspects of the paper that may require further clarification.

Thank you once again for your valuable feedback.

Dear Reviewer hpJw,

Thank you for your valuable time in reviewing our paper. We appreciate recognizing the importance of the problem we are investigating and the novelty of our analysis. We will fix the presentation issues in the final version.

In L209, the negative value is often understood to remove/erase the information from the residual stream. If this flexibility is removed from the model, it could affect the model's plasticity. It would be useful to see a more flexible gating function. The entropy can still be calculated using the absolute values of the coefficients (and does not need to change the model arch and training).

The reviewer suggests calculating entropy using the absolute value of coefficients without altering the architecture or training. We would like to clarify that we did not change any architecture or training, but all is calculated the same during the forward pass. We apply absolute value to the calculated value $f(x · K^T )$ during normal forward pass only to calculate knowledge entropy, where f is the non-linear activation function, SwiGLU. We believe this is the same calculation as the reviewer suggested. Please let us know if we missed something.

What does section 3.4 do to support the use of knowledge entropy? I understand the concept of attention entropy and next token entropy but how does that connect to the choice of using knowledge entropy? In fact, if all these entropy measures decrease over time, how can we justify that knowledge entropy is the main reason for the degraded plasticity even if it shows high correlation? Section 4.3 provides interventional evidence but similar intervention should be done with, say, attention entropy to ensure this is not the source for the lowered plasticity. I.e., if making attention entropy less peaky also increase retention and mitigates forgetting, then the role of knowledge entropy is not unique as other statistics reflect the same phenomenon. Would be nice to see more interventional evidence showing that knowledge entropy is the main contributor to the plasticity issue.

The reviewer raises questions about the connection between Section 3.4 and knowledge entropy and requests additional evidence to show that knowledge entropy, rather than attention entropy or next token entropy, is the primary factor contributing to reduced plasticity. We built the definition of knowledge entropy based on findings from [1]Geva et al. (2021), which demonstrated that the feed-forward layer captures knowledge. However, we acknowledge other works that explore the role of attention layers([2] Geva et al.(2023)) and knowledge circuits([3] Yao et al.(2024)). Thereby, in Section 3.4, we wanted to show those people who would be curious about the overall trend of the model apart from only the feed-forward layer that we observed that the overall model behavior exhibits similar trends across various entropy measures.

Nevertheless, we focused on defining knowledge entropy specifically in the feed-forward layer because our findings suggest it is the primary factor in degraded plasticity. Our experiments with resuscitation targeting attention layers showed some improvements in knowledge acquisition and retention, but these effects were less pronounced compared to resuscitation in the feed-forward layer (details of experiment setting and results in next paragraph). This suggests that, while attention entropy and next token entropy may reflect certain trends, knowledge entropy associated with the feed-forward layer plays a more significant role in plasticity degradation.

In the experiment where we introduce resuscitation into the attention layer, we apply temperature scaling to the dot product of key and query vectors, right before the softmax activation. While artificial interventions to increase Knowledge Entropy can be achieved by adjusting the parameters of the feed-forward layer at specific positions with low activation values, the same approach cannot be applied to alter Attention Entropy. This is because Attention Entropy reflects the probability distribution over prior token positions. To artificially increase Attention Entropy, we employed temperature scaling to reduce the sparsity of the probability distribution. We experimented with temperature values ranging from 1.5 to 3.0.

As we are not able to attach the figure, we add the figure in Appendix B.7 of the revised paper (Figure 17). The figure shows the results of resuscitation performed in the feed-forward and attention layers under varying temperature settings. The data indicates that resuscitation in the feed-forward layer consistently yields the highest knowledge acquisition rate and generally achieves the highest knowledge retention rate. This trend is reflected in both knowledge probe results and downstream performance. These findings suggest that Knowledge Entropy plays a critical role in influencing knowledge acquisition and forgetting, driving the observed differences in performance.

Dear Reviewer hpJw,

Thank you for your time and effort in reviewing our paper.

As the discussion period approaches its conclusion, we would like to kindly remind you of the upcoming deadline. We are happy to discuss any aspects of the paper that require further clarification or discussion.

Once again, we sincerely appreciate your valuable feedback and insights.

Dear Reviewer hpJw,

Thank you for your effort in reviewing our paper. We have addressed the notation overload issue you pointed out and corrected grammatical errors, revising the paper accordingly. Additionally, we have included an intuition about the relationship between decaying knowledge entropy and knowledge acquisition and retention performance in Appendix A.1 of the revised paper. We hope this revision could alleviate your concerns related to causality. It would be appreciated if we could discuss any remaining concerns during the extended review period.