Source-Aware Training Enables Knowledge Attribution in Language Models

Summary

This paper introduces a training recipe to elicit large language model to cite from pre-training sources. The proposed method starts with injecting document ids to a set of documents and continue pre-training on the documents with their ids, then perform instruction fine-tuning on set of prompts + document + document ids. The author validate their method in a carefully controlled synthetic setting and found that the proposed method effective improves attribution to pre-training data.

Rating

6

Confidence

3

Ethics flag

1

Reasons to accept

1. The study on pre-training source attribution is very timely — increasing the faithfulness of language model to parametric knowledge can make the output more verifiable and factual, leading to less hallucinations. 2. The proposed method is simple, does not induce too much overhead compared to pre-training (since it just continue pre-trains and do instruction SFT). 3. The experiment results shows the effectiveness of the method compared to the baseline, and the ablations on the location of the id should be of interest of many practitioners.

Reasons to reject

1. My first concern is that the author only conducted experiments in a **synthetic** setting and the construction of doc id relies on understanding the distribution of the documents, thus whether the method generalizes to real settings is questionable. Moreover, the author only conducted experiments on a small language model, whereas in larger models, it might be easier to make the model produce quotes from pre-training data just with lightweight preference optimization since a larger model memorizes better. 2. The controlled synthetic setting pre-trains a language model on a fixed set of high quality documents, but in reality, the documents are often very noisy. So whether the method generalizes to real language model pre-training is questionable. Given the above two points, I believe that this paper is novel and studies an important topic, but it is only limited to a very controlled setting.

Summary

In this work the authors propose the task of intrinsic source citation, namely associating pretraining documents with unique IDs that LLMs can be trained to reproduce when generating answers in a QA setting to attribute the source of the fact from parametric knowledge. This is a simple, clever, and innovative idea and the authors followed it through very thoroughly, with substantial grounding in existing literature and very thorough empirical results. I think it certainly could set other researchers off on a productive path trying related approaches. The paper is well-written and presents a good starting point for people interested in this area. I have one major issue, though - I'm very confused as to why the Doc IDs would include the names of the participants at all. You note this doesn't give away *all* the information, but why would we be providing the model with any hints whatsoever as to the correct Doc ID? Why is it not just a random hash of characters, or random sequence of unrelated words, or something? Isn't this making the task dramatically easier? Relatedly, you say each fact only appears in one document, but in data augmentation are the shuffled documents given the same Doc ID as the original? This is particularly concerning given the finding in Figure 4 that augmentation is strictly required for OOD attribution to work at all. The authors address this question explicitly with a subheading in section 5, "Is the model merely inferring the doc ID from the last name in the input?" But I don't feel that this paragraph is written sufficiently clearly, how exactly did you calculate this Jaccard Index? I'm not currently convinced by this. Even taking the findings at face value, I am genuinely questioning the long-term usefulness of the methodological approach presented, because even given the simplified, synthetic data you're working with performance drops off very sharply as documents get longer (Figure 5 right). Nevertheless if empirically valid this is still a useful negative result to present, in which case it's less a reason to reject this paper and more just illustrative of a weakness of the method. But my overall takeaway is I'm not confident we can trust this empirical evaluation and I really just am confused as to why this problem was introduced into these experiments at all. I can't find a justification for it in the paper, only some reasoning about why it's "not that bad" of a problem. But it just isn't necessary to have this problem, right? Am I missing something? I am willing to substantially revise this review if this can be explained in author response.

Rating

7

Confidence

3

Ethics flag

1

Reasons to accept

- Solid core idea - Thorough empirical tests - Well-written paper

Reasons to reject

- Potential serious problem with empirical soundness related to Doc IDs containing the names which are the core entity the questions are about - Unclear whether the approach presented will actually be at all effective with real-world data

Questions to authors

- In the introduction it's confusing to a reader because you talk about your approach as a post-pretraining recipe but then continually refer to enabling the model to attribute to documents seen in pretraining. - I would define "parametric knowledge" on the first instance because this term is relatively new, just a simple clause like "knowledge encoded in the model parameters" would help - In section 2, do we actually know RAG relies only on non-parametric knowledge? This doesn't seem necessarily so to me. - An example in section 4 gives a name as "Lara Jonhston", probably a typo for "Johnston". - Figure 3 needs a more clear legend, currently it has to be inferred from the x axis of the left panel. - Figure 5 appears before Figure 4; also in section 5.3 the textual references to figures are confusing and I think typos / bugs. - I am colorblind and find the figures marginally hard to read because many distinctions rely only on colors that are relatively close to one another; consider additionally using other features like point shape, line style, etc to make the distinctions more clear.

Summary

This work explores the important idea of parametric (i.e. only with model weights) attribution of knowledge in LLMs. The authors propose a “post-pretraining” method for tuning an existing model to 1. Memorize a corpus of information associated with unique doc ids and 2. Use that knowledge in a QA task, with citations to the underlying documents. They specifically apply this to a synthetic dataset of biological sketch information. The paper is reasonably clear and while the general idea of knowledge attribution is certainly extremely in fashion, this method is original (to my knowledge). I find this paper extremely interesting, well-motivated, timely and a natural idea! However, I have some concerns about the brittleness/effectiveness of this method and the underlying claims that this method could be used for general pretraining corpora, rather than a narrow and carefully groomed synthetic task. I do believe this paper is an important step in the correct direction.

Rating

6

Confidence

4

Ethics flag

1

Reasons to accept

This paper is clearly timely and has potential for great impact. LLMs capable of citing source documents would be extremely helpful at mitigating hallucinations and other LLM misbehavior. As the authors note “a significant portion of an LLM’s knowledge is acquired during pretraining”. This paper also conducts an interesting set of experiments on an admittedly artificial setting. The authors study sensitivity to issues like document ID placement, attention masking for packed batches, and effects of data augmentation. The idea is very natural and it is easy to see how some version of this method could be incorporated into full-scale LLM training.

Reasons to reject

I find the method and experiments in this paper extremely interesting and the authors make important contributions. However, I have concerns about effectiveness and connections to the claims made about the methods. Primarily, the abstract states “we explore source-aware-training… to teach the LLM to cite a supporting pretraining source when prompted. Source-aware training can easily be applied to pretrained LLMs off the shelf, and diverges minimally from existing pretraining/fine-tuning frameworks.” I do not believe the results here show that this method can “easily be applied” or that it diverges only minimally. The dataset is very synthetic, relying on what appear to be templatic, short, declarative sentences that uniquely express a single fact. The training instance is Table 1 is clearly an extreme divergence from the far messier presentation of facts in real pretraining data. Language in the intro and abstract suggest that this method will easily work on pretraining corpora, only addressing these concerns in the Limitations section (calling the method a proof-of-concept). Further, the abstract states “...we demonstrate that our training recipe can enable faithful attribution to the pretraining data without a substantial impact on the model’s quality compared to standard pretraining” - this claim is only referenced by perplexity analysis and I believe the results show the opposite of their claim. It is unclear how well the LLMs will work on any other tasks outside the very narrow biography QA setting. The authors use perplexity numbers to support the claim that their results maintain LLM quality. They measure wikitext-v2 perplexity as a measure of overall LLM quality, citing Radford et al. (GPT-2) as precedent. I understand that perplexities of different models and tokenization are not directly comparable, but we can observe relative trends. Figure 5 shows ppl going from ~19 to ~70, while Radford et al. shows ppl ranging from 29 to 18 at 117M parameters and 1.5B parameters respectively. Thus the increase in perplexity from source aware training seems to be FAR larger than the difference between a 100M model and a 1B parameter model (note the authors use a 1.1B model). Ultimately wikitext ppl seems a very weak measure of LLM quality and this method or the claims made about measurements should be revised. It also appears that Figure 3 shows several of the methods are substantially worse than the baseline of no document id QA (my apologies if I am misunderstanding the figure). **In conclusion, I believe this paper is very interesting but would strongly benefit from revising some language in the intro and abstract to more accurately reflect limitations of a synthetic study.**

Questions to authors

The section on attention masking is interesting and very important. However, the chosen method (“such that the ID tokens for a given document only attend to tokens from within that document.”) seems inverted - shouldn’t the mask prevent tokens from other documents from attending to the wrong document ID? Can you share any comments on randomly generated doc IDs, rather than using the truncated last name method? **Typos:** - “still substantially better compared REPEAT and is” - should say compared to repeat - Section 7 header misspells Conclusion - “Fig- ure 3 (right) shows perplexity trends during both pretraining and instruction tuning over BIOCITE and Figure 5 (Left) visualizes the tradeoff between LLM quality and OOD attri- bution.” **Should state Figure 5 instead of Figure 3 in the first clause?**

Summary

This paper proposes source-aware training, which trains the LLM to embed source document IDs in the pretraining corpus and instruction fine-tunes the model to learn to cite the supporting document in the generated response. Experiments on some synthetically-created data shows that source-aware training can enable faithful attribution to the pretraining data without a substantial impact on the model’s quality compared to standard pretraining.

Rating

4

Confidence

4

Ethics flag

1

Reasons to accept

Compared with retrieval-augmented generation (RAG), source-aware training allows the LLM to learn the attribution during the training, which avoids the difficulty of post-hoc retrieval process.

Reasons to reject

- However, the authors only evaluated the source-aware training on a synthetic pretraining corpus. It is hard to say whether the method could work when scaling up to real-world pertaining corpus. - We believe it is hard to scale up this method because it is costly to collect responses with their corresponding supporting documents. Compared with raw pretraining texts which can be easily collected on a large scale, (response, support document) pairs often do not naturally exist. Collecting such data can be time-consuming and costly. Compared with RAG which only requires a relatively small amount of (query, document) pairs to train the retriever, putting the attribution process at the pre-training stage would greatly increase the demand for such (not-easy-to-collect) data. - In source-aware training, document identifiers (usually numbers) are embedded into the model parameters. However, those document identifiers are specific to certain corpus. So different models might use different representations of document identifiers (e.g., one with a number and one with a random string). This causes problems when using multiple LLMs and it is less flexible than RAG because corpus-specific identifiers are directly stored in the model parameters. - This paper lacks a comprehensive comparison with the RAG models. It is not quite clear what's the benefit of source-aware training compared with RAG for generating responses with citations.

Typo correction

Dear reviewer, We would like to correct a typo in our rebuttal above in the RAG comparison bulletpoint. Here's the corrected version: * **Comparison with RAG:** The reviewer has made a good point that our approach can often be applied to scenarios where RAG is used. However, we would like to highlight that the goal of our work is fundamentally different from RAG: we aim to explore an intrinsic, model-based recipe for attribution rather than claiming superior performance to RAG, which is external to the model and requires integration with IR systems. Therefore, we do not compare to RAG or post-hoc techniques because the results of such comparison are **not** essential to our main point.

Response by Reviewer

I would like to thank the authors for providing a thoughtful response. I think the authors have addressed my concern on real world data, though I would prefer to see if only adding instruction SFT is able to let a pre-trained LLM output doc ids for a subset in its training data. However, I would like to maintain my score due to the two following concern: - As the model size gets larger and larger, we might not need this kind of effort to solicit the model to output doc ids, since it can memorize more. - I know that not many people have resources to train a 1.1B model, that is why I think more lightweight methods with just SFT or preference optimization on off-the-shelf LLM (e.g. LLama3) would be more impactful. As you have mentioned in your abstract, you introduced a "**post** pre-training recipe", which somewhat indicates that this can be applied to any **pre-trained** LM. Again, I think that the author did well in addressing some my concerns (synthetic v.s. real data), and I think the paper is novel and well written. But I still want to maintain my score due to the computational requirements and scalability of this method. It would make this paper much stronger to see if only adding doc ids to a few trusted sources and applying this method only during SFT to off-the shelf models.