Memorizing Documents with Guidance in Large Language Models

1 Introduction

Large language models (LLMs) have shown human-level performance on several tasks [Touvron et al., 2023; Brown et al., 2020]. The strength of LLMs comes from extensive model sizes and vast amounts of data. LLMs are trained with many documents in a corpus, and the parameters store information such as grammar, factual knowledge, and common sense [Geva et al., 2021]. Although an end-to-end training mechanism allows training from data, non-trivial memory location prevents tracing document contents.

A recent approach is inspecting activated units (neurons) in GPT that can reveal the semantic meanings, such as time, citation, and region [Bills et al., 2023]. However, problems such as spurious correlation and polysemantic occur with such a post-hoc analysis [Elhage et al., 2022]. A more trivial way is to store information in GPT memories with indexed memory locations; memory entries of documents are known in training. Such an architectural modification can trace the used memories in a document-wise manner. For this purpose, we propose document-wise memories, which have entries of memories for individual documents. Figure 1 shows the relationship between documents and documentwise memories. In the forward process, the document-wise entries guide the hidden representation to recall document contents. Document-wise entries could be either predefined or optimized. This work uses the second approach by utilizing guidance loss [Ho and Salimans, 2021].

We propose document guidance loss to (1) entangle memories and documents and (2) encourage different memory entries for documents. The original guidance [Ho and Salimans, 2021] increases the likelihood of conditional generation and decreases the unconditional part. The proposed loss modifies the unconditional part; we reduce the likelihood of document text with memories of other documents.

To map documents to memory entries, we use vector representations of documents; in short, DocRep and multi-layer perception (MLP) to map DocReps to memory entries. We study the relationship between DocReps and memory entries, assuming that close DocReps may have similar memory entries. Of course, trained document embedding [Li et al., 2021; Reimers and Gurevych, 2019] can motivate different memory entries as an inductive bias. However, we do not include the inductive bias on representation and initialize them randomly as MLP can reparameterize DocReps. We observe that more different memory entries are obtained with document guidance loss during training (Section 6.5).

This paper studies a link between document representation, memory entries, and the perplexity of document text with indexed memories. Figure 2 shows the perplexity of three documents with memories whose entries are generated from DocReps in 2D space. We use linear mapping from DocRep to memory entries. Therefore, the change in the DocRep space smoothly changes the memory entries, and consequently, the perplexity of document contents smoothly changes (the paraboloid shape). One way to explain the smooth perplexity change is the continuity assumption in metric spaces. The distance between two documents preserves the difference between memory entries; this aligns with the concept of Lipschitz continuity [Jones et al., 1993].

We compare memory entries with continuous and noncontinuous cases and empirically show the possibility of guidance loss with a linear function. In addition, we show that a nonlinear case does not work well with the proposed document guidance loss (Section 6.5).

Clarifying the knowledge location in LLMs is crucial for safe AI to protect user contents and believe the generated contents [Hacker et al., 2023]. This work contributes to the trustworthy AI communities by providing a study on designing document-wise memories, encouraging more reliable architectural and algorithmic design toward safe LLMs. We summarize our contribution as follows:

• We propose a document-wise memory mechanism to trace memory entries in memorizing documents.
• We propose document guidance loss to encourage different memory entries for documents and study the relationship between documents and memory space.

2 RelatedWork

This work studies document-wise neural memory for reliable LLMs. We review recent studies on (a) safety issues in LLMs, (b) memories in LLMs, and (c) memory structures.

2.1 Safety Issues in LLMs

LLMs show remarkable progress. However, several safety concerns arise [Hacker et al., 2023], including intelligence property infringement [Yu et al., 2023], hallucinations [Manakul et al., 2023], jailbreaks [Xie et al., 2023], machine generation detection [Mitchell et al., 2023], and privacy invasion [Pan et al., 2020]. This work proposes document-wise memories that can provide document-wise knowledge location as an approach to safe LLMs.

2.2 Knowledge in LLM

Identifying the knowledge location in LLMs is important to tackle safety problems. However, LLM is generally a black box, which is hardly explainable [Longo et al., 2024].

Recent work shows that lower layers have syntactic information while upper layers have semantic information [Geva et al., 2021]. Additional work shows that the neurons in LLMs are related to the factual knowledge [Dai et al., 2022]. Several studies investigated neurons in GPTs and highlighted the most activated concepts of neurons [Bills et al., 2023; Bricken et al., 2023]. Although inspecting knowledge location is interpretable, finding the knowledge location is not scalable for hyperscale LLMs. For example, GPT-3 has 175B parameters in 96 layers with embedding size 12,288 [Brown et al., 2020], and the post hoc interpretation of inspecting neurons requires extensive resources. In addition to that, spurious correlation can occur when interpreting neurons.

Recent studies on adaptations can provide separate knowledge locations by identifying adaptation processes. Common approaches are LoRA [Hu et al., 2022], editing knowledge neuron [Meng et al., 2022b], kNN memory injection [Khandelwal et al., 2020; Xu et al., 2023], and external knowledge adapation for LLMs [Diao et al., 2023]. These studies are promising approaches to utilize LLMs. However, more explainable and interpretable structures are required. For this purpose, we propose document-wise memories that can motivate reliable document-wise adaptations.

2.3 Neural Memory

A neural memory is a neural network structure that combines memories selected from memory entries. The memory entries, called keys, are obtained by multiplying an input and a key matrix [Sukhbaatar et al., 2015]. This structure promotes both the performance of LLMs and the interpretation of neurons. A recent study has shown that a transformer layer has a key-value neural memory architecture in a multi-layer perceptron (MLP) that stores semantic and syntactic meanings [Geva et al., 2021]. Numerous studies have verified knowledge in transformers: the existence of skill neurons for downstream tasks [Wang et al., 2022]; editing factual knowledge in transformers [Meng et al., 2022a]; kNN-based large memory design in internal layers of GPT [Wu et al., 2022]. Other fields of study include neural memories for Theory of Mind in reinforcement learning [Nguyen et al., 2023] and for anomaly detection in Vision tasks [Gong et al., 2019].

3 Document Conditional Generation

This section proposes document-wise memory architecture and document guidance loss. We consider a conditional generation with DocRep Ki of document Di for i ∈ {1, 2, · · · , N} where N is the number of documents. We use the term document Ki to indicate the representation of document Di for simplicity. For a passage (y1, y2, · · · , yt) in document Di where yt is the t-th token, the causal language modeling [Brown et al., 2020] has the following form

On the other hand, the conditional generation of the passage with DocRep Ki is the following form

Our goal is to make document-wise memory entries for the conditional generation. One possible approach is the allocation of disjoint memories for documents, which is not scalable as the number of memories is proportional to the number of documents. Instead, we train a memory selection, allowing overlapping of memory entries. 1

3.1 Document-Wise Memory

Document-wise memory has the MLP architecture in a transformer, and memories are selected conditionally on the DocReps. Consider an MLP module in a transformer layer

where x is input vector, K, V are key and value matrices with activation σ and biases bk, bv respectively. We denote σ(K1x + b1) by KeyTok(x). DocRep K generates memory entries KeyDoc(K) to block entries softly. The documentwise memory is the following form

where the first key KeyTok(x) is the selection of memories from the hidden representation x in the forward of language models. In contrast, the second key KeyDoc(K) is the memory entries for document K and is generated by a function g. We interpret KeyDoc(K) as the memory selection of document K. Figure 3 shows the document-wise memory architecture. To entangle document contents and memories selected from DocReps, we introduce document guidance loss.

1 Learning to encode information in fixed vector dimensions is a typical property of neural networks [Elhage et al., 2022].

3.2 Document Guidance Loss

Ensuring the conditional generation in Equation 2 can be trained with classifier guidance [Dhariwal and Nichol, 2021] or classifier-free guidance [Ho and Salimans, 2021; Nichol et al., 2022]. We propose a document guidance loss based on the classifier-free guidance. Consider passage y in document D whose DocRep is K. For an LLM with parameter θ, the implicit classifier considers the following equation

To increase the likelihood of DocRep K, the numerator part must be increased while the denominator part decreases. This is proportional to the following equation

where α controls the ratio between two probabilities. To encourage different memories for documents, we decrease the likelihood of the passage conditionally on DocRep K-(̸= K). This forgetting process decreases the likelihood (reciprocal of energy [Du and Mordatch, 2019]) of text in document K with another condition K-. Equation 6 becomes

Although Equation 7 can ensure the desired properties, the negative part is unbounded; that is, the cross entropy loss ranges from (-∞,∞). To stabilize the training, we assume that the conditional generation with the negative DocRep has a low likelihood Plow. Finally, we have

and the loss L is the following form

where τ is a constant, LCE(ˆ yK, y) is the cross entropy loss of passage y and conditional generation ˆ yK ∼ Pθ(·|·,K). The right hand side is the guidance loss part, which encourages forgetting y from the memories of negative DocRep K-.

3.3 Negative DocReps for Guidance Loss

We use the term positive document for a document that includes passage y, and negative documents that are assumed to have low perplexity for the passage. Negative DocReps are essential to encourage different memories. We suggest three possible choices of negative DocReps: zero, other, and random.

• zero: K- = 0.
• other: K- = Kj such that Kj ̸= K, j ∈ [N].
• random: K- = R where Ri ∼ N(0, ϵ).

For random, ϵ covers the document representation space. The memories selected from negative DocReps increase the perplexity of the text, and only positive DocReps can encourage low perplexity. The choice of negative DocReps affects the overall knowledge structure in memories.

4 Metric Spaces

We formulate memory selections from documents with metric spaces. Consider three metrics (a) dDoc, (b) dDocRep, (c) dKey for spaces (a) documents, (b) DocReps, and (c) memory entries (keys), respectively, and two functions f, g which map documents to representations, and consecutively memory selections respectively. Figure 4 shows the relationship between three metric spaces. This section analyzes the memory selection of function g with a continuity assumption.

4.1 Continuity Assumption

When g selects memories continuously, two close DocReps will have similar memory selections. When g is τ-Lipschitz, we obtain the following bound.

Proposition 1 (Lipschitz Continuity for Memory Selection). Let K1,K2 be two DocReps with dDocRep(K1,K2) ≤ ϵ. When g is τ-Lipschitz, dKey(g(K1), g(K2)) ≤ τϵ.

Proposition 1 shows that the memory selection difference is bounded by the factor ϵ, which is the difference between DocReps. In other words, for two DocReps bounded by ϵ, the memory selection difference could not be more than τϵ. Proposition 1 explains the paraboloid shape in Figure 2. As g is a linear function in the example, the memory section is smooth, and the perplexity is also smoothly changed.

However, smoothness could hurt performance for randomly initialized DocReps. Consider two documents K1 and K2 with different contents. When DocReps are closely initialized, the memory entries are similar, too. If we want different memory entries for these documents, one feasible approach is finding a nonlinear ˆ g which holds dKey(ˆ g(K1), g(K2)) > τ ϵ condition. However, we observe that the nonlinear function has pitfalls when trained with document guidance loss.

4.2 Caveats of Non-Liptschitz Continuity

We train document guidance loss with nonlinear g. Figure 5 shows the perplexity of the nonlinear function g with ReLU activation for three documents with the zero negative DocRep. A three-layer MLP has higher perplexity only on the zero. On the other hand, two-layer MLP has higher perplexity locally around zero. This observation reveals that continuity assumption affects local regions, and deeper layer depth does not differentiate memories properly.

Training document-wise memories involves two learnings: memory selection and memorization. Memory selection with guidance loss can provide different memory entries when the DocRep space is continuously linked. We conjecture that a smooth memory selection manifold with document guidance loss encourages different entries. However, if the selection is nonlinear, there could be pitfalls with document guidance loss. A more constructive hypothesis and experiments are required to study nonlinear cases (see also Section 6.5).

5 Experiments

We train Pythia 1B [Biderman et al., 2023] to memorize Wikitext-103-v1 [Merity et al., 2017] by replacing the last MLP with the proposed MLPdoc whose memory size is 128. We individually train document-wise memories for 10, 20, and 50 documents with guidance α = 0.1 and τ = 2.5. For baselines, we train two types of memory modules without guidance. Shared is the MLP in Equation 3, and Add is a module that directly adds differential memory entries. We also evaluate three activation types for document memory entries: ReLU, Tanh, and Sigmoid, which affect memory selections. We make the source code publicly available. 2

6 Results

6.1 Activation Types for Key Generation

We show how activation types affect memory selections in guidance loss. Figure 6a and 6b show the guidance loss for 10 and 50 documents with different activations. The Sigmoid function shows the fastest training speed, followed by Tanh, GeLU, and ReLU. However, the selected memories of Sigmoid are not visually document-wise entries. Figure 7 shows the memory selection of the activations. The memory entries of Sigmoid are similar for documents compared to Tanh, known as a gating mechanism, and ReLU, known to cut decisions. Visually, Tanh and ReLU are proper inductive biases for constructing document-wise memory entries.

2 https://github.com/fxnnxc/DocGuidanceLLM

6.2 Comparison with Baselines

Figure 8a shows the cross-entropy loss of Shared, Add, and guidance with Tanh activation, guidance α = 0.5, and zero negative DocRep. The proposed method shows slower learning than other methods because the guidance loss has a negative part that hinders memorization. This observation is supported by ablation on α in Figure 8b. When α is large, the forgetting dominates the training.

Although Shared is better at memorizing, the generated text is a mixture of several contents. Table 1 shows the generation results with the prompt Wikipedia. When document memories are mixed (Shared), the next word prediction is the most likely word under all documents. On the other hand, document-wise memories do not mix all the memories and provide more document-related content.

6.3 Different Memories for Documents

The goal of the guidance is to have different memory entries. We show the perplexity of three documents with memories selected from all DocReps individually in 2-dimensional space. Figure 9 shows the perplexity of three documents with DocReps in the space. We observe that Sigmoid activation shows smooth perplexity over the space, while Tanh shows visually different DocRep regions for the perplexity of three documents. In addition, Tanh shows a sharp increase in perplexity compared to Sigmoid. We observed the same pattern for ReLU. Therefore, ReLU and Tanh are the proper choices for document-wise memories with document guidance loss.

We quantitatively evaluate this observation with ROUGE scores [Lin, 2004]. The memories selected by DocReps can be evaluated in two ways. First, the conditional generation must include the document contents, and second, the conditional generation must not include the other document contents. To evaluate the second, we report IV-ROUGE (Inverse of ROUGE score) defined by

where Dj is the jth document, and Gi is the generated contents only with the memories of ith document. The IVROUGE score is low when memories do not contain the contents of other documents. We train document-wise memories with DocReps of 32 dimensions for three seeds and generate 32 tokens for six simple prompts, such as He is, with document memories. Tables 2 and 3 show the ROUGE and IV-ROUGE scores, respectively. ReLU and Tanh are consistently better at document conditional generation in both metrics. We believe localized knowledge (high ROUGE and low IV-ROUGE scores) can enhance knowledge editing [Yao et al., 2023] by editing disentangled document memories.

6.4 Negative DocReps

Table 4 shows the ROUGE precision score for 20 documents and three types of negative DocReps. The zero DocRep shows the best ROUGE score, and the other case shows the best IV-ROUGE score. The ROUGE scores correlate with the amount of forgetting. The random DocRep encourages forgetting for most memories and provides the lowest ROUGE score. Similarly, the other DocRep includes more forgetting than the zero. Similarly, other case removes contents from memories of other documents and provides the lowest IVROUGE score for both activations.

The guidance loss part also supports these metrics (Figures 10a and 10b). Note that the random and other negative DocReps have higher losses than the zero negative DocReps. The random case did not show convergence to zero as sampling includes the positive DocRep.

6.5 Nonlinear Memory Selection

Linear memory selection entangles documents and memories with document guidance loss. However, we observed that nonlinear memory selections do not work well. We compare 1 (linear) to 4 layers with ReLU internal and final activations, adding MLPDoc (128 memories) after the last decoder layer. The models are trained with 10 documents, α = 1.0, τ = 4.5, and random negative DocReps. Figure 12 shows the final memory entries of the first five documents. We observed high similarity in nonlinear memory selections and more dead neurons (not activated) as the number of layers increases.

To quantitatively evaluate memory selections, we train linear and nonlinear cases for five seeds and measure L0, which is the number of non-zero entries, L2 norm, and pairwise distance, which is ||g(Ki) - g(Kj)||2 for two DocReps Ki and Kj , and normalized pairwise distance by L2 norm of memory entries. All metrics are averaged over ten documents (Figure 11). The number of nonzero entries (L0) increases for all cases, meaning the rank of memory usage increases. However, the normalized pairwise distance decreases as the layer depth increases. The linear memory provides more different memory entries than the nonlinear cases. The large gap between linear and nonlinear indicates the limitation of guidance loss with nonlinear cases.

The benefit of linear memory selection can be explained by the continuity of memory selections. We conjecture that the guidance loss, a maximization and minimization game by selecting different memories, works well with continuous functions as most of the points in the manifold are connected. This claim is experimentally supported by memory selection (Figure 12), perplexity heat map (Figure 9), and pairwise distance (Figure 11). However, we do not suggest any explanation for the worse performance of nonlinear memory selections as the exact behavior of nonlinear is unclear. We believe more structural hypotheses and experiments are required for nonlinear cases, which are currently out of scope.

7 Discussions

This work introduces document-wise memories with minimal assumption on document representations by randomly initializing them. In many cases, random DocReps with the Lipschitz continuous function are insufficient. As noted, nonlinear memory selection currently has limitations with the proposed guidance loss. We suggest two approaches: 1) inspect nonlinear memory selections and solve the problem, or 2) utilize high quality text embedding [Lee et al., 2024] for the linear memory selection.

This work focuses on storing document contents. However, the implications for the downstream tasks are not presented. We believe document-wise memories can benefit downstream tasks where the factuality of documents is necessary, such as medicals [Sallam et al., 2023; Dave et al., 2023]. In addition, this paper does not tackle a large number of documents. A hierarchical document structure may work better than directly applying guidance loss. This work aims to provide reliable LLMs with known document entries. Therefore, we encourage interpretable entries even for a large number of documents. Lastly, we acknowledge that false positive memory entries (e.g., unrelated document contents in memories) can exist. Therefore, additional studies are necessary to verify the semantic meaning of trained memories.

8 Conclusion

Storing documents in the traceable locations of LLMs is a crucial research topic. This paper studies document-wise memories in LLMs. We propose document guidance loss to entangle document contents and document memories, encouraging different entries for documents. We also provide a theoretical view of memory selections with metric spaces and continuity assumptions. The experimental results show that the proposed guidance loss provides different memory entries with linear memory selection while leaving nonlinear memory selection as an open problem.

Ethical Statement

The proposed document-wise memories can provide traceable representations for data, promoting transparent and reliable language models.

Acknowledgements

This work was partly supported by Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. 2022-000984, Development of Artificial Intelligence Technology for Personalized Plug-and-Play Explanation and Verification of Explanation; No. 2019-0-00075, Artificial Intelligence Graduate School Program (KAIST); No. 2022-0-00184, Development and Study of AI Technologies to Inexpensively Conform to Evolving Policy on Ethics), and Samsung Electronics MX Division.