Source-Aware Training Enables Knowledge Attribution in Language Models
Authors: Muhammad Khalifa
David Wadden
Emma Strubell
Honglak Lee
Lu Wang
Iz Beltagy
Hao Peng
Venue: COLM (2024)
Keywords: attribution, pretraining, citation generation, factuality
TL;DR: We study training language models to provide support for their parametric knowledge by citing the pretraining data.

Abstract: Large language models (LLMs) learn a vast amount of knowledge during pretraining, but they are often oblivious to the source(s) of such knowledge. We investigate the problem of intrinsic source citation, where LLMs are required to cite the pretraining source supporting a generated response. Intrinsic source citation can enhance LLM transparency, interpretability, and verifiability. To give LLMs such ability,
we explore source-aware training---a recipe that involves (i) training the LLM to associate unique source document identifiers with the knowledge in each document, followed by (ii) an instruction-tuning stage to teach the LLM to cite a supporting pretraining source when prompted. Source-aware training borrows from existing pretraining/fine-tuning frameworks and requires minimal changes to the model architecture or implementation. Through experiments on synthetic data, we demonstrate that our training recipe can enable faithful attribution to the pretraining data without a substantial impact on the model's perplexity compared to standard pretraining. Our findings also highlight the importance of pretraining data augmentation in achieving attribution.

Published as a conference paper at COLM 2024
Source-Aware Training Enables Knowledge Attribution in Language Models
Muhammad Khalifa**, David Wadden Emma Strubell
Honglak Lee", Lu Wang", Iz Beltagy",, Hao Pengll*
University of Michigan", Allen Institute for AI Carnegie Mellon University
University of Illinois Urbana-Champaignl khalifam@umico am@umi ch. edu
Abstract
Large language models (LLMs) learn vast amount of knowledge during pretraining, but they are often oblivious the source(s) of such knowl- edge. We investigate the problem of intrinsic source citation, where LLMs are required to cite the pretraining source supporting a generated response. In- trinsic source citation verifiability. To give LLMs such ability, we explore source-atoare training- - recipe that involves (i) training the LLM to associate unique source doc- ument identifiers with the knowledge in each document, followed by (ii) instruction-tune tunning stage to teach the LLM to cite supporting pretrain- ing source when prompted. Source-aware training borrows from existing pretraining/ fine-tuning frameworks and requires minimal changes to the model architecture or implementation. Through experiments on synthetic data, we demonstrate that our training recipe can enable faithful attribu- tion to the pretraining data without substantial impact on the model's s perplexity compared to standard pretraining. Our findings also highlight the importance of pretraining data augmentation in achieving attribution]
enhance LLM transparency, interpretability, and
Introduction
Large language models (LLMs) often generate con- tent that is not based on factual information ([i et al. 2023 Ye et al. 2023a). As LLMs are pre- trained over noisy web data that often contains in- accurate or outdated content, users should be able to verify LLM outputs by checking their sources. Moreover, concerns about copyright infringement (Min et al. 2023; Longpre
Ethan Howke and Jule Delpy
the movie Before Survise?
Knoeledga
al., 2023), privacy Violations Kim et al. 2024), data contamination (Shi et al. 2023), and toxic content (Gehman al.) 2020) in LLMs emphasize the need for techniques to identify and trace the origins of information included in models' responses. It is therefore de- sirable if LLMs can provide supporting evidence for their responses by citing or attributing the outputs to the sources they draw upon (Rashkin et al. 2023 Huang & Chang 2023 Li et al. 2023b). Beyond improving the models' transparency, attribution allows for deeper understanding of the relationship between training data and model behaviors, thereby offering a pathway to refine the quality of pretraining data.
Phetesning
prosuiing
Langaage
Figure 1: Intrinsic source citation: The lan- guage model cites pretraining document(s) from which acquired its relevant paramet- ric knowledge.
We focus on intrinsic source citation, where the LLM should cite source documents from the pretraining data from which it acquired its relevant parametric knowledge. Compared to
"Wor partially done while these authors were at the Allen Institute for AL Code and data available here: http:/ccc...ccc//c//////ttttttat

---

Published as a conference paper at COLM 2024
post-hoc techniques (He et al. 2023 Gao et al. 2023a), intrinsic source citation is inherently tied to the model itself, enabling more taithful attribution to its parametric knowledge, and opening up unique opportunities for improved interpretability (Alvarez Melis & Jaakkola 2018 Marasovic et al. 2022).
retrieval-based approaches such RAG (Lewis et al. 2020 Guu et al. 2020
To this end, we explore source-aware training to enable an LLM to cite the pretraining source supporting its parametric knowledge. Our motivation is three-fold. First, significant portion of an LLM's knowledge is acquired during pretraining, therefore citing evidence for this parametric knowledge can greatly enhance the LLM trustworthiness. Second, the standard practice for LLM pretraining neglects attribution, which explains why LLMs currently fail to provide reliable citations (Agrawal et al. 2023 Zuccon et al., 2023). We aim to explore training procedure that naturally facilitates citation of the pretraining data. Finally, from a scientific perspective, it is intriguing to investigate whether and how current language models can be trained to reference their pretraining data.
We inquire: Given an off off-the-shelf LLM, can we train it to attribute its generations to the supporting sources from the pretraining data? Our goal is to cite the pretraining documents themselves (see Figure []. Our setup mirrors existing frameworks for LLM (continual) pretraining and can be summarized
follows: We take off-the-shelf LLM, continue to pretrain it on a corpus associating each document with a unique identifier, then fine- tune it to answer questions about the acquired knowledge while providing citations. The citation is achieved by generating an identifier of document supporting the answer. Continual pretraining done as in prior work, with the main difference of injecting the document identifiers into the pretraining data danimml changes in the model's architecture or implementation are needed.
To study the generalization over this task and simulate a realistic fine-tuning setting, we limit our instruction tuning stage a subset of the pretraining documents (in-domain) and evaluate the model's attribution ability over the remaining (out-of-domain) documents. We run experiments over synthetic pretraining corpus of fictitious biographies and show that LLMs can achieve reasonable attribution when answering question about the out-of- domain documents.
Our contributions are summarized as follows:
To the best of our knowledge, this work is the first to study intrinsic source citation and investigate the ability of current LLMs to cite the source of their parametric knowledge.
We explore source-aware training, which we apply to off-the-shelf LLMs to give them the ability to attribute their outputs to the pretraining sources. On synthetic data, we show that such training can achieve reasonable attribution without strongly hurting the LLM perplexity compared to standard pretraining.
We examine the impact of various training strategies on attribution, such as training to cite the gold document (Section 3.3) and data augmentation (Section 5.3) and our findings can inform future efforts to train attribution-capable models at a large scale.
2 Related Work
Language Model Attribution. Attribution is gaining more attention recently as inter- pretability and grounding of language models become increasingly important. Generally speaking, approaches to achieve attribution can be classified as either retrieval-based or model-based. Retrieval-based approaches include retrieval augmentation (RAG) (Lewis) et al. |020 Guu et al. 2020 Borgeaud et al. 2022 Izacard et al. 2023) and host-hoc attribu- tion (He et al. 2023 Gao et al. 2023a). RAG approaches enable attribution by providing retrieved context for the LM to use, and teaching LM how to cite the retrieved context (Nakano et al. 2021 Menick et al. 2022). The major limitations of RAG approaches are the the retrieved data for generation (Petroni et al., 20200 Li et al. 2002a),, and that they only work on non-parametric knowledge. Post-hoc approaches (He et al. 2023
lack of guarantee that the model is relying
Gao et al.] 2023a) attribute the LM outputs by retrieving the sup-
2

---

Published as a conference paper at COLM 2024
porting evidence given the model's response, but have been shown to produce non-accurate citations (Liu et al. 2023).
Model-based techniques involve prompting the model directly to generate citations for its parametric knowledge (Weller et al. 2023; Zuccon et al. 2023) or scaling techniques such as influence functions (Koh & Liang 2017) to large models (Grosse et al. 2023). Model- based attribution is arguably more faithful than retrieval-based approaches as the citation mechanism is intrinsic to the model (Alvarez Melis & Jaakkola 2018 Marasovic et al. 2022). However, standard approaches to pretraining LMs do not take into account the need for the language model cite its pretraining data, which is where our work comes into play. Bohnet et al. (2022) proposed the task of attributed question-answering and evaluated the attribution performance different systems using the AutoATS metric (Rashkin et al. 2023 Gao et al. 2023a). In addition, they fine-tuned PaLM (Chowdhery et (023)) to generate both an answer and URL pointing to Wikipedia page supporting the answer in generative
retrieval style (lay et al. 2022: Wang et al., 2022 Although this setup is similar to ours in that we require the LM to generate the document identifier as well, their setup is basically a variation of RAG where the LM acts as the retriever.
Citation Generation. There is a large body of work on the task of citation generation in the scientific domain, where the goal is to cite an appropriate article given particular context (McNee et al., 2002 Nallapati et al. 2008) or to generate text citing one article in relation to another (Xing et al. (0200] Luu et al. 2020). relevant work to ours is Galactica (Taylor et al. 2022), which leverages the underlying citation graph in the pretraining data to learn to predict citations given a context. Notably, Galactica is trained to leverage citations of scientific articles in the pretraining data, while our work explores citation of all the pretraining documents, extending beyond scientific articles. Gao et al. (2023b) introduced benchmark for the automatic evaluation of LM citations and|Ye et al (2023b) proposed a method to improve language model grounding by fine-tuning the language model on responses that are well supported by their citations. However, their setup is restricted to citation of retrieved rather than parametric knowledge.
Generative Retrieval. Our work is somewhat related to generative retrieval, where an auto regressive model is trained to act retriever (Wang et al. 2022 Tay et al. 2022). Generative retrieval typically relies on transformer model to map given query
information retrieval (IR) system
document identifier that is likely contain an answer to the query. While our task also requires the language model to generate a document identifier, we differ from generative retrieval in at least two ways. First, our goal is to generate an identifier pointing to document containing the already generated answer rather than a document that likely to contain the answer. Second, generative retrieval merely learns mapping from query to document identifiers, while our setup is concerned with both acquiring knowledge via the nort prediction objective over the documents and associating acquired knowledge with its source.
Source-Aware Training
Our training framework is designed to easily integrate with existing pretraining pipelines. We minimize its deviations from established pretraining practice, and involves almost no modifications to the model architecture or implementation. Each document in the pretraining corpus is assigned a unique document identifier (ID) and our goal to train a language model that can respond to user prompts by providing both response and an ID referring to the source document of the model's knowledge.
Setup. Our evaluation follows the attributed question answering setup (Bohnet et al. 2022), where given an input prompt z, the LLM output will consist of a tuple (r,c) where r is the response (e.g., the answer to a question) and c is the identifier of the document in the pretraining data that supports the answer. Following standard LLM training setups, our recipe has two stages: Pretraining (Section 3.1) and instruction tuning (Section 3.2). Instruction tuning trains the model to be able to attribute the generated responses to
3

---

Published as a conference paper at COLM 2024
Pretraining with injection
Pretraining corpus
In-domain docs
domain docs
Attribution evaluation
Instruction tuning
support your enower.
Answer
does Kylc
Wher Kyle Study?
A: Woshingtor State Univerity <Ldsbro goo bus-buus-dea jase-jn///dd
Language Model
A: Columbu, Ohio <idbrod-buss -asbroo--ba---o---ollee
containe
the fact "Kyle Goodain Fors in Columbux Ohio
Figure 2: Training and evaluation setup: The pretraining corpus is split into in-domain and out-of- domain documents. The in-domain documents used to create instruction tuning examples, and the out-of-domain documents are used for attribution evaluation.
supporting documents it has seen during pretraining, The pretraining stage will involve all documents by nature, but the instruction tuning step restricted to subset of the pretraining documents. This restriction is due to the potential cost of curating instruction tuning data from all the pretraining documents, that is in addition to the training overhead incurred by instruction tuning (Zhou et al. 2024).
After training, we measure out-of-domain (OOD) attribution: whether the model can attribute knowledge to documents that are only included in the pretraining data but not in the instruction tuning data. We therefore split the pretraining corpus into in-domain and OOD subsets. The in-domain data is used to create instruction training examples, while the OOD documents are used for evaluation, as shown in Figure)
3.1 Pretraining with Doc ID Injection
Doc ID injection. The pretraining phase has two goals: (i) memorizing knowledge via next-word prediction (same as established LLM pretraining), and (ii) associating knowledge within source document with its ID to enable OOD attribution. We aim to achieve the second goal by injecting the document TD into the document before training. An important consideration is the location and frequency of injecting the document ID.
Formally, given pretraining corpus of documents (,(1),x(2) ( (m)) and their correspond- ing IDs
where each x00 is sequence of tokens x(() (xf²
(i)
(i) ... |x
x2
and each c(i) (c))
is sequence of tokens of its identifier. Our pre- training aims to learn the language model parameters that maximize the objective maxg E,, log P()(0);)) where g(i) is the ID-injected version of the document x(0), We inject the doe into document a with different strategies, each of which corresponds different &I Particularly, we experiment with the following strategies:
" ,...... (cl
(1) (....
1. NO-ID: Standard pretraining without ID injection: 2 :
2. DOC-BEGIN: Inject the ID once before the first token in the document: f (2,,Cr,.., C2, Clcl1x1, X2, X }x
We omit the superscript for brevity.

---

Published as a conference paper at COLM 2024
Hayden Bush taught various schools before joining Fox Corporation content developer. Josiyn Valdez received her undergraduate degree Forensic Science from Harvard University. Ellie Kemp began her career Apple before joining Gracle in Cary, North Carolina. sid-bur-kem-valee kam-val-/ii>>
Answer the question and support your
A: Where did Joslyn Valciez study?
Ellie Kemp began her career Apple before joining Cracle Cary, North Caro lina <idb bus ked-
B: Hervard Universily
Chain-af-thought attribution: generate anewer followed by the rest of the document then the doc ID
Pretraining document with ID injected
Figure 3: Example of chain-of- ththugghtt attribution. During both training and inference, the model cites the remaining part the document before generating the doc ID.
3. DOC-END: Inject once after the last token in the document. This is equivalent to
8 := X2, ..., xz1= C1, C2...,
4. REPEAT: Inject the ID after every sentence in both in-domain and OOD documents.
Here, := CI, Maj) Xsy, C2, C Gi] Xsy. C1, 91j where Xs, are the tokens corresponding to the j-th sentence document and assuming x has sentences.
Disabling cross-document attention. To maximize GPU utilization during continual pretraining, the typical practice packs several pretraining documents within single training sequence separated by the end-of-sentence <eos> token. As a result, if cross-document attention is enabled with causal masking, the doc ID tokens for certain document will attend to preceding tokens from other documents and vice versa. Our initial experiments showed that this severely hurts attribution, since the model will associate the doc ID of given document with knowledge from other documents in the same training sequence. To avoid this, we disable cross-document attention in all our experiments.
3.2 Instruction Tuning
In addition to pretraining, we further adapt the model to (i) recall the appropriate knowledge as a response to the prompt and (ii) cite the ID of the document supporting the response This stage does not teach the model any new knowledge, but merely aims at eliciting memorization of both the knowledge and doc ID by instruction tuning (Wei et al.) 2021 Zhang et al. 2023). Given I examples, the i-th example is a tuple (2(),,"1, (())) where 2!! is the prompt (instruction + query), p(ū is ground-truth response, and c(i) is the ID of a document that supports the response. The model is trained with the objective max _n log (2()) (()())
The instruction-tuning examples only come from the in-domain documents, and use the simple instruction "Answer the following question and provide evidence for your answer." Figure 2 shows a fine-tuning example from BIOCITE. During the standard LLM pretraining, i.e., with NO-ID, we remove the doc ID part from instruction tuning examples. Following Taylor et al. (2022), we surround document IDs with two learned special tokens <id> and </id> during both pretraining and fine-tuning.
3.3
Chain-of-Thought Attribution
The setup in the previous section trains the model to recall the doc ID immediately after the answer. Another setup we explore is where the model is asked cite the ground-truth document (or part of it) before generating the doc ID. This can be regarded as an instance of
chain-of-thor tf-thughht (CoT) (yee et al., 2021 Wei et al. 2022 In CoT, we inject the document using DOC-END, then the model is trained to cite the rest of the document after the answer up till the doc ID. Figure33 shows an example of the CoT setup.
3DOCEEND results in the same training objective as in DSI TTay et al. 2022), where the model is trained to generate the ID given the full document. While this objective was shown to work for the information retrieval setup, we find that fails generalize attribution.
"The instruction tuning examples are curated from the pretraining data such that for given prompt, we already have the reference document the model should cite. More details in Section 4
5

---

Published as a conference paper at COLM 2024
Document: Marleigh Austin works SpaceX. Marleigh Austin studied at the University of Arkansas, Fayetteville. Isaiah Brown studied Graphic Design. Isaiah Brown was born on October 19, 1930. Lora Johnston was born on May 30, 1989. Lora Johnston works at Microsoft Teams. Kyle Goodwin studied Washington State University. Kyle Goodwin works Campari Group.
Doc ID: auto-goo-aus-J0
Instruction tuning example:
Q: Where does Lora Johnston work?
A: Microsoft Teams ## <id>bro-goo -oobbroooooo-ss--ohh/sss
Table An example document in BIOCITE with its unique identifier and example question. Documents in BIOCITE are constructed by sampling biographies of fake individuals (highlighted in color) from BioS (Zhu & Li 2023) and then sampling several facts from each biography.
4 Data
To have a controlled experimental setting, we rely on pretraining knowledge in the form of atomic synthetic facts. We now describe how construct BIOCITE- synthetic pretraining corpus.
Sampling facts. BIOCITE is based on the BioS dataset (Zhu Li 2023), a collection of biographies of fake people where each biography lists six different facts about each person: birthdate, birth city, study major, university, employer, and work city Each attribute is described using a corresponding template. For example, the birth city fact is described by "<person name> was born in <birth city>." " To avoid co-reference issues when sampling facts, the person's full name is mentioned in all the facts.
To simulate realistic pretraining data that often include facts about different entities, construct each document in BIOCITE as a collection of facts from at least two different biographies in BioS. More specifically, to construct one document d, we first sample the number of biographies ny ~ Uniformi [2.... NMMxNBBil) Then, we sample ny biographies from BioS without replacement. Finally, we sample random number of facts from each one in the n, biographies and combine these to form the document. We allow the same combination of biographies to create document only once and allow each fact to appear only once in BIOCITE In our experiments, we generate 100K documents total using MMaxNBo 4.
Creating questions. The input prompts for BIOCITE will take the form of factoid questions about the different facts such as "Where does Lora Jonhston Work?". Question generation is done by mapping each fact in the document to corresponding question template. For example, a fact about a person's 's birth city mapped to the question "Where was <full name> born?"
Doc ID format. It has been shown that the document ID design plays a role in generative retrieval performance (Tay al. 2022 Pradeep et al. 2023 Sun et al. 2024 and we observed the same during our initial experiments. When designing a doc ID, we need to be careful not to make the task too easy, where the model can infer the doc ID from the input question without actually performing attribution, but also not too hard, where no semantic overlap exists between the doc ID and the content of the document. Using the last names of the individuals included in document fits these criteria for two reasons. First, two facts from the same person will most likely exist in different documents. Second, the same last name can be shared by different biographies, whose individuals differ only in the first name. That means relying on the last name will not sufficient to trivially predict the correct doc ID.
Details about reproducing BioS are in Appendix A.1
In this work, we assume each fact in BIOCITE is mentioned in exactly one document and leave the extension of this work to multi-doc citation to future work.

---

Published as a conference paper at COLM 2024
We choose to use dash-separated concatenation of the 3-letter prefixes of the last names from the biographies that make up the document, shuffled randomly. Table 1 shows an example document, its ID, and question extracted from it. Exact dataset statistics are in Table in the Appendix.
Data augmentation. LMs struggle to generalize at knowledge extraction over OOD docu- ments (i.e., document that were not seen during fine-tuning) without sufficient amount of redundancy where the LM will be exposed to the same fact in different formats/ postsiossssonsss (Zhu & Li 2023 Allen-Zhu & Li 2023 Berglund et al. (023). In large-scale pretraining se- tups, this is achieved by scaling the pretraining data but as we study attribution on a smaller scale, we achieve the same effect of redundancy via data augmentation. We mainly apply doc-level augmentation, by shuffling the sentences in each document Naug 3 times, where Naug is the number of augmentation samples. Unless otherwise stated, our experiments will include document-level augmentation of the pretraining data, and we will explore the effect of augmentation on attribution in Section 5.3
Experiments and Results
5.1 Experimental Details
The pretraining corpus is split 50-50 into in-domain and OOD subsets, respectively. Training is done over 80% of the in-domain question, and we show performance in the remaining 20K. OOD evaluation is performed over 20K questions randomly sampled from the OOD documents. The QA performance is evaluated using the token exact match (EM) with the gold answer. During inference, we prompt the model and let it generate a response first, then append the special token <id> and continue decoding until the model generates the </id> token. We use constrained beam search Cao et al. 20 Tay et al. (2022) to force the model to generate doc IDs that appeared in the pretraining data.
We evaluate attribution by measuring whether the cited document supports the question- answer pair. Precisely, we measure the gold document ID recall over cases where the answer is correct, where recall is evaluated using Hits@k with k E [1, which measures whether the gold ID is in the top beams. To monitor the impact of our attribution training on the model quality, we monitor the perplexity over Wikitext-v2 (Merity et al. 2017) during training, as done in previous work (Radford et al. 2019). The model we use for all experiments is Tiny 1.1B (Zhang et al. 2024) which we pretrain for 10 epochs with learning rate x 10-5 and iustruction-tuning for epochs with a learning rate of 1 x 10 5 During both pretraining and fine-tuning, we apply a linear decay scheduler and use batch size of 128, weight decay of 0.02, and learning rate warm-up of one epoch.
5.2 Results
Downstream QA Performance. We start by evaluating the QA performance on OOD quuestions. Figure 4 (left) shows answer match over BIOCITE with different document ID injection strategies. The model can achieve OOD answer match > 80%, showing that the model has well memorized the pretraining knowledge. We also note that DOC-BEGIN achieves much worse QA performance than other strategies, and we hypothesize that DOC- BEGIN conditions the model to expect the ID when citing knowledge, causing a mismatch during inference when the ID is absent.
OOD attribution depends on ID injection strategy. The ID injection strategy plays a major role in OOD attribution achieved by source-aware training. As shown in Figure (right), placing the ID only once with DOC-BEGIN or DOC-END performs poorly. We hypothesize that both cases train the model to associate the full document-rathe than individual facts- with the document ID. Precisely, DOC-END conditions the model on the full document when generating the doc ID, but the ev evaluation requires the model to predict the ID given
hhu.ginnggae..////mmyym///////aa (B-innerrmddiatee annterrediate-ste---eee stepee-step11111111 1431k-31
7

---

Published as a conference paper at COLM 2024
In-domain
Out-of-Domain
0.6
0.6
Indoomain
0.8
Nrower EM
hits@@
0.4
0.4
0.2
0.2
0.2
8.0
0.0
begin
end
CoT
repeat
1
end
instruction tuning epochs
doc
strategy
Figure 4: Left: Answer EM over questions from in-domain and OOD documents after fine-tuning epoch with different ID injection strategies (3.ction3.1).. The LLM can generalize well to out-of-domain (OOD) questions in all document ID locations, although both in-domain and OOD answer EM scores degrade with DOC-BEGIN. Right: Hits@1 over in-domain and OOD questions during instruction tuning. Only REPEAT and DOC-END CoT can achieve OOD attribution.
individual facts not full documents. This is an instance of LLM generalization failures in knowledge extraction discussed in prior work (Zhu & Li 2023: Allen-Zhu & Li 2023) and explains why REPEAT is substantially better, since it trains the model to predict the ID after each fact, making it easier for the model to associate individual facts with the ID.
Chain-of-thought attribution helps. REPEAT may be unfavorable, since the number of
pretraining tokens will noticeably increase, bringing additional training overhead 8 Besides, the model quality will be negatively impacted since document IDs are not natural text, which is reflected in the perplexity over Wikitext-v2 shown in Figure 4 (right). The question
here is whether source-aware training can yield OOD attribution while injecting the doc ID once. Interestingly, the chain-of-the chain-of-thought setup (Section 3.3) achieves reasonable OOD attribution without requiring repeating the doc within the document. however, that the CoT setup adds extra training and inference overhead required to generate chain part of the output. Another interesting observation is that REPEAT and DOC-END + CoT achieve better OOD answer EM compared to NO-ID (e.g., 88.8% with REPEAT 80.9% with NO-ID). We conjecture that source-aware training improves the model grounding to the pretraining data, which reflects the QA performance.
worth noting,
The results above suggest that source-aware training can teach the model to attribute its parametric knowledge to their pretraining sources, with one key choice to consider: the doc ID injection strategy. Another key component is document augmentation, which we discuss in the next section.
5.3 Additional Analysis
Impact on perplexity. Now we study the impact of different document ID injection strate- gies on the LLM quality measured in terms of perplexity over Wikitext-v2. Figure (right) shows perplexity trends during both pretraining and instruction tuning over BIOCITE and Figure (Left) visualizes the tradeoff between LLM quality and OOD attribution.
(right)
First, we note that perplexity increases during training in all setups due to the domain shift incurred by training BIOCITE, which does not resemble real text. We can use the perplexity with NO-ID baseline and observe how other setups compare to it.
As expected, REPEAT exhibits the worst perplexity, since frequent ID injection means training on more non-natural text. We also note that DOC-BEGIN shows very high perplexity even though the doc ID is injected once, showing that best to include the doc ID later rather than earlier in the document. Finally, even though DOC-END CoT leads to worse perplexity than NO-ID, it is still substantially better compared REPEAT and Pareto-optimal as shown in Figure (Left). These results that DOC-END CoT strikes the best balance between OOD attribution and maintaining the model's quality.
8The number of pretraining tokens increased by about 80% with REPEAT compared NO-ID.
8

---

Published as a conference paper at COLM 2024
Mikltxt2
attribution quality
391
dec if bege
repert cna
Log Likelloood (willtext-/2)
perplesity
-42
-43
pretraining Rnowning
-45
38
04
training epochs
Hits@1
4facts
Figure 5: Left: LLM quality vs. OOD attribution. Higher better for Hits@l and Log Likelihood. Optimal is top-right corner. DOC-END + CoT is Pareto-optimal as it strikes the best balance between LLM quality and OOD attribution. Right: OOD attribution performance with different gold document lengths.
OOD attribution vs. document complexity. We analyze how OOD attribution varies with the complexity of the document measured in terms of the number of facts when training with REPEAT and DOC-END + COT. In ( we plot OOD attribution measured with Hits@10 10 changes as the number of facts in the gold document changes. We observe consistent trend where documents with more facts are harder to cite. This can be explained by the limited representational capacity of the doc IDs: Documents with more facts require the doc TD be associated with more knowledge.
Impact of Document Augmentation. We com- pare two types of data augmentation methods: doc- ument and fact augmentation, and the goal is to assess which type of augmentation necessary for OOD attribution. Document augmentation is done by permuting the facts within a document Naug times and is what our experiment so far have re- lied on. Fact augmentation duplicates the facts in document in Naug different random documents. Figure 6 shows OÔD answer match and Hits@1 as Naug varied and where Naug 1 means no augmentation. While answer match improves us- ing fact-level augmentation, Hits@1 remains the same and only improves when we apply document augmentation. Document augmentation appears necessary for the model to associate the doc ID with the facts in the document.
000 antaee natch
0O0 14tx511
1.0
1.0
0.8
0.8
0.4
0.2
0.2
0.0
NaIz
Stc sugmentat.
fact augmantation
Figure 6: The effect document- and fact- level augmentation on OOD answer match and Hits@1. Naug = means no aug- mentation applied. We show results when training with the doc ID injection strategy REPEAT.
Random doc IDs. To study the effect the doc ID de- sign, we experiments with randomly generated doc
EM
Hits@1
IDs, where each ID is an 8-digit string (e.g., 53476421). We compare the performance of random IDs to our last name- based design when using REPEAT. Table T shows the out-of- domain answer exact match and hits@1 BIOCITE. While the model still acheives reasonable attribution with random IDs, they underperform our truncated last name design since it is harder for the model to associate numerical to- kens (e.g., the token "3") than natural language tokens (e.g., "rog") with document facts.
Last names Random ID
88.8
47.3
95.9
39.8
Table 2: Out-of-domain QA and attribution performance with dif- ferent document ID designs.
Is the model merely inferring the doc ID from the last name in the input? Since the doc IDs are constructed as a concatenation of the first three letters of last names in the facts in the documents, the LLM could cheat by predicting doc IDs that contain the prefix of

---

Published as a conference paper at COLM 2024
Q: What university did Adelyn West attend?
Q: what city does Alissa West work?
Answer:
University of Pittsburgh.
Answer: New Orleans.
Top predicted ids:
Top predicted Ids:
jen-lyn-wes
wes-gri
jen-wes-bur
wes-mcc
jen-cob
wes-wat-vau
Gold document: Adelyn West was born
Gold document: Angelina Grimes was born on December 27, 1916. Angelina Grimes studied at...
August 7, 1954. Alissa West lives in New Orleans. Adelyn West studied at University of Pittsburgh....
Table 3: Two out-of-domain input questions with the same last name and the model predictions. The model predicts totally different doe IDs showing that the model relying on the fact content itself rather than merely predicting IDs that contain the last name prefix.
the last name in the question. If that were the case, then given two questions about two individuals with the same last name (e.g,, Q1: when was Maria Thomson born?; Q2: where does Jake Thomson live?), their generated doc IDs would significantly overlap. We measure this by predicting the top 10 doc IDs for every out-of-domail question. Then for every pair of questions about individuals sharing the same last name (there are 352K such pairs), we compute the overlap between the top 10 predicted doc IDs corresponding to each question using Jaccard index. The average Jaccard over all such pairs is only 0.08, indicating that this is not the case. Table shows two examples of such outputs and the top three predicted doc IDs for each question.
Limitations
Our work presents a proof-of-concept (PoC) on source-aware training and, as with all PoCs, it has limitations:
We rely on synthetic rather than real-world data, and the main motivation for this is control for potential confounding factors introduced by using real data, and which might indirectly affect attribution. Another limitation is that we restrict the form of knowledge be attributed to factual world knowledge, which we particularly choose since the utility of supporting factual knowledge is more obvious compared to other types of knowledge such as commonsense knowledge, for example.
Synthetic data:
Small-scale experimentation: Our experiments are done using a relatively small pre- training corpus and model size. This is mainly due to the massive compute that would be required to run hundreds of experiments using a billion-scale pretraining corpus. Nonetheless, believe the insights revealed by our experiments are valuable and can benefit future research involving large-scale experiments.
Cost of source-aware training: Our experiments show that due to inherent limitations with LLMs, generalization which may practically increase the cost of pretraining. One workaround is to realize that not all pretraining data should be cited. For instance, we could select sources that we know to be reliable (e.g., Wikipedia) and only apply source-aware training to these.
anat-of-oomin--documee
requires data augmentation,
7 Conslusion
In this work, we studied intrinsic source citation, where models are required to provide support for their parametric knowledge by citing evidence from the pretraining data. We explored a two-stage training process that involved training with document identifiers injected into the pretraining data and then instruction tuning. Our results showed that such source-aware training can enable parametric knowledge attribution in language models when data augmentation is applied. We believe our findings are valuable for future research training verifiable and trustworthy models.
10

---

Published as a conference paper at COLM 2024
References
Ayush Agrawal, Lester Mackey, and Adam Tauman Kalai. Do language models know when they're re hallucinating references? CoRR, abs/ 18248, 2023. doi: 10.48550/ AR0IV.230..... 18248. URL otct:p:://oii..rg/1.....44 wwt..////100..0/////k.k.... 2305. 18248
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.2, knowledge doi: 10.48550/ AR0I0.X3XXXXX11040000000 URL
manipulation. CoRR, abs/2309.14402, 2023.
https://ddi.org: /ttgg//..85500000 www.50///rXii...
14402
David Alvarez Melis and Tommi Jaakkola. Towards robust interpretability with self- explaining neural networks. Advances in neural information processing systems, 31, 2018.
Lukas Berglund, Meg Tong, Max Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Korbak, and Owain Evans. The reversal curse: Lims trained on" a is b" fail to learn" is a". Xiv preprint ar xrXiv:2309112888 2023.
Bernd Bohnet, Vinh Q Tran, Pat Verga, Roee Aharoni, Daniel Andor, Livio Baldini Soares, Massimiliano Ciaramita, Jacob Eisenstein, Kuzman Ganchev, Jonathan Herzig, et al. Attributed question answering: Evaluation and modeling for attributed large language models. arv preprint ar arxi:22212.088: 2022.
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 2206- 2240. PMLR, 2022. URL htttps://prrceeedinggg. mlr.
borgeaud22.. potpeeaa///..////rrrraaaaaaa..a. html
Nicola De Cao, Gautier Izacard, Sebastian Riedel, and Fabio Petroni. Autoregressive entity retrieval. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https:///penreerier net/ #orumiid==22Fe6UUS3V
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Luyu Gao, Zhuyun Dai, Panupong Pasupat, Anthony Chen, Arun Tejasvi Chaganty, Yicheng Fan, Vincent Y. Zhao, Ni Lao, Hongrae Lee, Da-Cheng Juan, and Kelvin Guu. RARR: researching and revising what language models say, using language models. In Anna Rogers, Jordan L. Boyd-Graber, and Naoaki Okazaki (eds.), Proceed- ings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2023, Toronto, Canada, July 9-14, 2023, pp. 16477-16508. Association for Computational Linguistics, 2023a. doi: 10.18653 /W1/2023AACL-LONG.CCCC https://doi.co org/ 10. 18653/v1/2022 acl-long....... 910
URL
Tianyu Gao, Howard Yen, Jiatong Yu, and Danqi Chen. Enabling large language models to generate text with citations. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceed- ings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6010, 2023, pp. 6465-6488. Association for Computational Linguistics, 2023b.
btttpp://aacanthhotn wtwphooooyyoooooooooooo 2023. emnlp mommp-mmai.... ..i..398nn
Samuel Gehman, Suchin Gururangan, Maarten Sap, Yejin Choi, and Noah A Smith. Realtox- icityprompts: Evaluating neural toxic degeneration in language models. ar preprint ar 0Xie:2009.11462. 2020.
11

---

Published as a conference paper at COLM 2024
Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, et al. Studying large language model generalization with influence functions. arXiv preprint arXio:2308.032:
2023.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In International conference on machine learning, pp. 3929-3938, PMLR, 2020.
Hangfeng He, Hongming Zhang, and Dan Roth. Rethinking with retrieval: Faithful large language model inference. CoRR, abs/ 2202200000000 2023. doi: 10.48550/ A205///AXXX.23000000000000 URL bittps://dol wtt.:oor//111.0...11 arXiv. 10300.000000
Jie Huang and Kevin Chen-Chuan Chang- Citation: key to building responsible and accountable large language models. CoRR, abs 2023, doi: 10.48550/ARXI 2307.02188 URL https::///o.. http.///////...///r/....... 2307 02185
Gautier Izacard, Patrick S. H. Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi-Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: Few-shot learning with retrieval augmented language models. J. Mach. Learn. Res., 24:
251:1-251:43, 2023.
URL http://jpmm- hwtpaiirrrrgppppppappaeeee pttpppapers//22223333333 0037 .htl
Ziwei Ji, Nayeon Lee, Rita Frieske, Tiezheng Yu, Dan Su, Yan Xu, Etsuko Ishüi, Ye Jin Bang, Andrea Madotto, and Pascale Fung. Survey of hallucination in natural language generation. ACM Computing Surveys, 55(12):1-38, 2023.
Siwon Kim, Sangdoo Yun, Hwaran Lee, Martin Gubri, Sungroh Yoon, and Seong Joon Oh. Propile: Probing privacy leakage in large language models. Advances in Neural Information Processing Systems, 36, 2024.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pp. 1885-1894. PMLR, 2017.
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval- augmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33:9459 9474, 2020.
Daliang Li, Ankit Singh Rawat, Manzil Zaheer, Xin Wang, Michal Lukasik, Andreas Veit, Felix X. Yu, and Sanjiv Kumar. Large language models with controllable working memory. In Anna Rogers, Jordan L. Boyd-Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguistics: ACL 2023, Toronto, Canada, July 9-14, 2023, pp. 1774-1793. Association for Computational Linguistics, 2023a. doi: (2.18633/V12200 HINDINGS-ALLIDE UtL/httts://ooi.orr/// 1tt.//000..633//1/0/0000 www.:o/g...iiiiinn-aaaaee .12
Dongfang Li, Zetian Sun, Xinshuo Hu, Zhenyu Liu, Ziyang Chen, Baotian Hu, Aiguo Wu, and Min Zhang. A survey of large language models attribution. arXiv preprint arXiv:2311.03731, 2023b.
Nelson Liu, Tianyi Zhang, and Percy Liang. Evaluating verifiability in generative search engines. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, Singapore, December 6-10, 2023, pp. 7001-7025. Association for Computational Linguistics, 2023. URL https:://claatholloog org/2023
findings emnlp.40.. 467
Shayne Longpre, Robert Mahari, Anthony Chen, Naana Obeng-Marnu, Damien Sileo, William Brannon, Niklas Muennighoff, Nathan Khazam, Jad Kabbara, Kartik Perisetla, et al. The data provenance initiative: A large scale audit of dataset licensing & attribution in ai. Xiv preprint arXiv:2310.16787, 2023.
Kelvin Luu, Rik Koncel-Kedziorski, Kyle Lo, Isabel Cachola, and Noah A. Smith. Citation text generation. ArXio, abs/2002.00317, 2020. URL https: :ttpps://app //api. semanti cscholar org/ CorpusID:211010022
12

---

Published as a conference paper at COLM 2024
Ana Marasovic, Iz Beltagy, Doug Downey, and Matthew E. Peters. Few-shot self- rationalization with natural language prompts. In Marine Carpuat, Matie-Catherine de Marneffe, and Iván Vladimir Meza Ruíz (eds.), Findings of the Association for Com- putational Linguistics: NAACL 2022, Seattle, WA, United States, July 10-15, 2022, pp. 410-424. Association for Computational Linguistics, 2022. doi: 20111663/V1/02//
FINDINGS-NAACL.L. URL https:///o// http:://doooogg///// wtt.......//.../....... 10....11653//1100 worp22...iiniii findings -aacl
31
Sean M McNee, Istvan Albert, Dan Cosley, Prateep Gopalkrishnan, Shyong K Lam, Al Ma- munur Rashid, Joseph A Konstan, and John Riedl. On the recommending of citations for research papers. In Proceedings of the 2002 ACM conference on Computer supported cooperative work, pp. 116-125, 2002.
Jacob Menick, Maja Trebacz, Vladimir Mikulik, John Aslanides, Francis Song, Martin Chadwick, Mia Glaese, Susannah Young, Lucy Campbell-Gilinghamm, Geoffrey Irving, et al. Teaching language models to support answers with verified quotes. ar Xiv preprint arXiv:2203.1112 11144, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL
itttps://openrreviee www.seewwweeffoooeee met//orumiidd-y///...
Sewon Min, Suchin Gururangan, Eric Wallace, Hannaneh Hajishirzi, Noah A Smith, and Luke Zettlemoyer. Silo language models: Isolating legal risk in a nonparametric datastore. arXiv preprint arXiv:2308.04433 2023.
Reiichiro Nakano, Jacob Hilton, Balaji, Jeff Wu, Long Ouyang, Christina Kim, Christo- pher Hesse, Shantanu Jain, Vineet Kosaraju, William Saunders, Xu Jiang, Karl Cobbe, Tyna Eloundou, Gretchen Krueger, Kevin Button, Matthew Knight, Benjamin Chess, and John Schulman. Webgpt: Browser-assisted question-answerrigg with human feedback. CoRR, abs/2112.09332 2021. URL http:p/://aarc....gg//a/////....
Ramesh M Nallapati, Amr Ahmed, Eric P Xing, and William W Cohen. Joint latent topic models for text and citations. In Proceedings of the 14th ACM SIGKDD international confer- ence on Knowledge discovery and data mining, pp. 542-550, 2008
Maxwell I. Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models. CoRR, abs /2/1/221100011,, 2021. URL https://arxiv
lrrg/abs/2122. 2/12/000/1112
Fabio Petroni, Patrick S. H. Lewis, Aleksandra Piktus, Tim Rocktäschel, Yuxiang Wu, Alexan- der H. Miller, and Sebastian Riedel. How context affects language models' factual pre- dictions. Dipanjan Das, Hannaneh Hajishirzi, Andrew McCallum, and Sameer Singh (eds.), Conference on Automated Knowledge Base Construction, AKBC 2020, Viriual, June 22-24, 2020, 2020. doi: 10.24432/ 24432/C5201W.
URL http:ttst//:i.orrgggge
24432/C5201V
Ronak Pradeep, Kai Hui, Jai Gupta, Adam D Lelkes, Honglei Zhuang, Jimmy Lin, Donald Metzler, and Vinh Q Tran. How does generative retrieval scale to millions of passages? arXiv preprint ar wix::2305.118444 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. 2019.
Hannah Rashkin, Vitaly Nikolaev, Matthew Lamm, Lora Aroyo, Michael Collins, Dipanjan Das, Slav Petrov, Gaurav Singh Tomar, Iulia Turc, and David Reitter. Measuring attribution in natural language generation models. Computational Linguistics, pp. 1-66, 2023.
Weijia Shi, Anirudh Ajith, Mengzhou Xia, Yangsibo Huang, Daogao Liu, Terra Blevins, Danqi Chen, and Luke Zettlemoyer. Detecting pretraining data from large language models. arXio preprint arXi::2310.1678, 2023.
13

---

Published as a conference paper at COLM 2024
Weiwei Sun, Lingyong Yan, Zheng Chen, Shuaiqiang Wang, Haichao Zhu, Pengjie Ren, Zhumin Chen, Dawei Yin, Maarten Rijke, and Zhaochun Ren. Learning to tokenize for generative retrieval. Adoances in Neural Information Processing Systems, 36, 2024.
Yi Tay, Vinh Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Prakash Gupta, Tal Schuster, William W. Cohen, and Donald Metzler. Transformer memory as a differentiable search index. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Ad- vances in Neural Information Processing Systems 35: Annual Conference on Neural Informa- tion Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 De- cember 9, 2022, 2022. URL http:://paperre nips. cc/ wwwppper..ppp..pcoo pappraapppppppe fileeell///eee. (alles/paper/20e
htmm
Ross Taylor, Marcin Kardas, Guillem Cucurull, Thomas Scialom, Anthony Hartshorn, Elvis Saravia, Andrew Poulton, Viktor Kerkez, and Robert Stojnic. Galactica: A large language 2022.
model for science. arXiv preprini arXio:2211.0900
Yujing Wang, Yingyan Hou, Haonan Wang, Ziming Miao, Shibin Wu, Qi Chen, Yuqing Xia, Chengmin Chi, Guoshuai Zhao, Zheng Liu, al. A neural corpus indexer for document retrieval. Advances in Neural Information Processing Systems, 35:25600-2555 226614 2022.
Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. Xiv preprint ar Xrv:22109.016522 2021.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V
Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824-24837, 2022.
Orion Weller, Marc Marone, Nathaniel Weir, Dawn Lawrie, Daniel Khashabi, and Benjamin Van Durme. according to..." " prompting language models improves quoting from pre-training data. Xiv preprint ar Xix::2230 13252, 2023.
Xinyu Xing, Xiaosheng Fan, and Xiaojun Wan. Automatic generation of citation texts in scholarly papers: A pilot study. In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel Tetreault (eds.), Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 6181-6190, Online, July 2020. Association for Computational Linguistics. URLhttpss http:///aaaantlooyyo.o wrg/2020.0000. 2020. acl mmain
doi: 10.
550
Hongbin Ye, Tong Liu, Aijia Zhang, Wei Hua, and Weiqiang Jia. Cognitive mirage: A review of hallucinations in large language models. CoRR, 2bs/2309,06794, 2023a. doi: 10.48550/ ARXIV.2309 067944. UtLpstss::///ot org/10. 48550/arXiv http:://oo......0000666666/..... 2309. 06794
Xi Ye, Ruoxi Sun, Sercan Ö Arik, and Tomas Pfister. Effective large language model adapta- tion for improved grounding. arxiv preprint ar Xiv::333 099533,, 2023b.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. Xiv preprint ar XrXiv:2401.0233 2024.
Shengyu Zhang, Linfeng Dong, Xiaoya Li, Sen Zhang, Xiaofei Sun, Shuhe Wang, Jiwei Li, Runyi Hu, Tianwei Zhang, Fei Wu, et al. Instruction tuning for large language models: A survey. Xiv preprint ar XiX:v:230 10098, 2023.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, et al. Lima: Less is more for alignment. Advances in Neural Information Processing Systems, 36, 2024.
Zeyuan Allen Zhu and Yuanzhi Li. Physics of language models: Part 3.1, knowledge storage and extraction. CoRR, /abs/2309 14316, 2023. doi: 10.48550/ 10.48550/ARXT/22000 ARXTV URL
ftttp:://doi.org//11 48100.8855/aaaa arXiv. 2309. 14316
Guido Zuccon, Bevan Koopman, and Razia Shaik. Chatgpt hallucinates when attributing answers. In Proceedings of the Annual International ACM SIGIR Conference on Research and Development in Information Retrieval in the Asia Pacific Region, pp. 46-51, 2023.
14

---

Published as a conference paper at COLM 2024
A Data details
A.1 Reproducing BioS
As Zhu & Li (2033 has not been publicly released, we reproduced our own version as follows. Each biography lists six different facts about each person: birthdate, birth city, study major, university, employer and work city. To fill values for each fact, we start by compiling list of first names, last names, company names, cities, universities, and majors. Examples and counts of the seed entities are shown in Table
Then we construct a single biography by sampling a unique full name (first and last names could repeat) and then to make up each fact, we sample random suitable entity. For example, for birth and work cities, we sample a random city and for employers, we sample company name and so on. For birthdates, we generate random date in the range: January 1st, 1900 - December 31st, 2099. Each fact is constructed via a corresponding template, shown in Table 5 To construct questions for instruction tuning with BIOCITE, each attribute is mapped to corresponding template question as shown in Table
Entity
Count
Examples
First Name
851
John, Mary, David, Aria...
Last Name
537
Smith, Johnson, Williams...
Company Name
284
Apple, Google, Microsoft....
City
199
New York, Chicago, Tampa....
University
178
University Alabama, University of Michigan,
Major
107
Industrial Design, Fashion Design, Psychology...
Table 4: Counts of the compiled entities used to reproduce BioS Zhu & Li 2023
Attribute
Fact template
Birthdate
<full name> was born <birthdate> <full name> works <company> <full name> lives in <city> <full name> was born in <city>
Work place
Work city
Birth city
University <full name> studied at <university> <full name> studied <major>
Major
Table 5: Templates used to construct different biograohy facts for BioS.
Attribute
Question template
Birthdate
When was <full name> born? Where does <full name> work?
Work place
Work city
Where does <full name> live? Where was <full name> born? Where did <full name> study?
Birth city
University
Major
What did <full name> study?
Table 6: Templates used to construct questions from each document in BIOCITE.
15

---

Published as a conference paper at COLM 2024
A.2 Dataset Statistics
Table 7 shows statistics of training and instruction tuning data for both BIOCITE and Wikipedia.
Size
Pretraining
#documents
100K
#facts/ sentss
408K
#tokens
5.7M
avg. sents per doc avg. tokens per doc
4.1
56.9
Instruction tuning
#examples #tokens
186K
3.1M
Table 7: Statistics for the BIOCITE before data augmentation. After data augmentation with Naug the number of pretraining tokens goes up 17.1M.
16


Reviews:
[{"review_id": "fCxeDEN1fy", "reviewer": "Reviewer_AHVg", "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.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments on our paper! We aim to address your comments below:\n\n* **Synthetic data:** We conduct our study on synthetic data because we want to control for different confounding factors that could impact attribution, e.g., complexity, noise in real pretraining data, etc. We believe our findings are valuable for large-scale pretraining efforts and **reviewer Szuz** agrees that “it is easy to see how some version of this method could be incorporated into full-scale LLM training”. \n* **Difficulty collecting SFT data:** In fact, we believe it is fairly easy to collect such data. We first note that the SFT data is only collected over a subset (e.g., 30%) of the pretraining corpus, enabling easy the scaling of our recipe to realistic pretraining setups. Also, this data collection can be automated akin to recent works that employ instruction-tuned LLMs to generate pretraining/SFT data. \n* **Specificity of doc IDS:** In practice, one can associate each model with its own id->document mapping, built from its training data, and then it should be possible to map back to the original document. Do you have a specific use case where this remains an issue?\n* **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 essential to our main point.\n\nTo alleviate concerns about the applicability of our approach to real data, we have applied source-aware training to 50K docs from [CCNews](https://huggingface.co/datasets/vblagoje/cc_news). For instruction tuning, we prompt ChatGPT to generate question-answer pairs. We used the exact same training hyperparameters as in the paper. Here are the results: \n\n| Doc ID strategy | EM| Hits@1 |\n|-|------|--|\n| ID-BEGIN | 21.0 | 12.5|\n| ID-END| 23.3 | 12.1 |\n| ID-REPEAT | 27.8 | 36.2|\n\nWe find the same trend as with BioCite: ID-REPEAT performs best and our approach achieves reasonable OOD attribution on a pretraining corpus of real data. \n\n\nWe hope we have addressed your concerns, and we hope you can reconsider your score. Please let us know of any more concerns!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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. \n- 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. \n- 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. \n- 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.", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "65NaKvnF5E", "reviewer": "Reviewer_Szuz", "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. \n\nThe 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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for your positive feedback on our paper and recognition of the potential impact of our submission! We aim to address your concerns below:\n\n\n* **Language in the intro:** We will revise our language in the intro and abstract to reflect that our results were done on a synthetic setup. \n* **Perplexity analysis:** We believe that _domain shift_, rather than degraded model quality, is the primary reason for the perplexity increase on WikiText. This is due to training on a specialized corpus such as BioCite. Therefore, in Figure 5, we can see that perplexity still increases even with standard no-id pretraining. We argue that it is unfair to compare our perplexity numbers to Radford et al. as their general-purpose pretraining does not cause such domain shift. We believe that scaling our recipe to large-scale, general-purpose data can mitigate this. We will include this discussion, and revise our language to be more specific to perplexity rather than model quality in general. \n* **QA performance:** Figure 3 (left) shows QA performance with different doc ID injection methods. On the contrary, we see better QA performance with our source-aware training compared to standard pretraining + SFT. We discuss this observation in section 5.2 (under chain-of-thought attribution)\n* **Attention masking**: The mask should prevent document tokens from attending to wrong doc IDs and doc ID tokens from attending to wrong document tokens. For example, assume a sequence of two documents that looks like this ($T_{ij}$ is token $j$ for doc $i$ and $U_{ij}$ is doc ID token $j$ for document $i$): [$T_{11},T_{12},T_{13},U_{11},T_{21},T_{22},T_{23},T_{24},U_{21}$]\n\nThe attention mask will prevent $U_{21}$ from attending to $T_{11},..T_{13}$ and $T_{21},...,T_{24}$ from attending to $U_{11}$. We will revise this part to make both scenarios clear. \n\n* **Randomly generated doc IDs**: We initially experimented with randomly generated doc IDs, but found them to underperform our truncated last name design since it is harder for the model to associate numerical tokens (e.g., the token ‘3’) than natural language tokens (e.g., ‘rog’) with document facts. The table below shows the performance of each:\n|| OOD Answer EM | OOD Hits@1 |\n|-|-|-|\n| Last-names truncated  | 88.8%| 47.3%|\n| Random ID| 95.9%| 39.8%|\n\nWe will include results with random IDs in future revisions.\n\nWe hope we have addressed your concerns, and please let us know if you have any more questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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.\n\nPrimarily, 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). \n\nFurther, 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.\n\nIt 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).\n\n**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?\n\nCan you share any comments on randomly generated doc IDs, rather than using the truncated last name method? \n\n**Typos:**\n\n- “still substantially better compared REPEAT and is” - should say compared to repeat\n- Section 7 header misspells Conclusion\n- “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?**", "extra_scores": {}}, {"review_id": "H3Rj083Ygf", "reviewer": "Reviewer_6yNB", "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. \n\nThis 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.\n\nI 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?\n\nRelatedly, 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.\n\nThe 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.\n\nEven 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.\n\nBut 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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": true, "final_justification": "", "rebuttal": "Thank you so much for your feedback. We are glad you found the paper idea clever and innovative! We address your concerns below:\n\n* **Doc ID design:** We note that our approach could work with any doc ID format. We mainly use truncated last names design since it performs better and converges faster than e.g. a random ID, while not making the task trivial. Below we compare our doc ID design with random IDs under the ID-REPEAT strategy: \n\n|| OOD Answer EM | OOD Hits@1 |\n|--|-|-|\n| Last names truncated | 88.8%| 47.3%  |\n| Random | 95.9% | 39.8%|\n\nWe observe that our approach still works reasonably well with random numerical IDs.\n\n* **Data augmentation:** We intend to say that each fact is associated with a single doc ID and will revise our wordings to make that clear. We admit that the requirement of data augmentation is a limitation, but as we suggest, it can be mitigated by selectively augmenting target subsets of the pretraining data that we wish the model to cite. \n* **Jaccard index:** We want to verify the model is not merely predicting doc IDs containing prefixes of the last name in the input question. If that were the case, then given two questions with the same last name, their generated doc IDs would significantly overlap. We measure this by predicting the top 10 doc IDs for every out-of-domain question. Then for every pair of questions (there's 352K such pairs) about individuals sharing the same last name (e.g., Q1: when was Maria Thomson born?; Q2: where does Jake Thomson live?), we compute the overlap (using Jaccard) between the top 10 predicted doc IDs corresponding to each question. The average Jaccard over all such pairs is only 0.08, indicating that this is not an issue. We will elaborate on this in the revision.\n* **Doc length effect:** While attribution is harder as documents become longer since more facts need to be associated with the same doc ID, we speculate that this effect may be mitigated if a larger model is used. We will investigate the effect of model size on attribution in the revision to validate this.\n\n> i really just am confused as to why this problem was introduced\n\nWe assume you refer to the doc ID problem addressed above. If not, please let us know.\n\n* **Results on real data:** Kindly refer to our response to **reviewer a1pn**, containing new results of applying our approach to CCNews—a corpus of real data.\n\nWe hope we have addressed your concerns, and please let us know if there's anything else you’d like us to clarify!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- Solid core idea\n- Thorough empirical tests\n- 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\n- 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.\n- 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\n- In section 2, do we actually know RAG relies only on non-parametric knowledge? This doesn't seem necessarily so to me.\n- An example in section 4 gives a name as \"Lara Jonhston\", probably a typo for \"Johnston\".\n- Figure 3 needs a more clear legend, currently it has to be inferred from the x axis of the left panel.\n- Figure 5 appears before Figure 4; also in section 5.3 the textual references to figures are confusing and I think typos / bugs.\n- 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.", "extra_scores": {}}, {"review_id": "K11qhUpikk", "reviewer": "Reviewer_a1pn", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for acknowledging the novelty and importance of our paper! We will address your concerns below:\n\n* **Synthetic setup:** Our motivation for using synthetic data is to control for different confounding variables that may affect attribution, leading us to the wrong conclusions. Our controlled setup allows us to take a step towards tackling this novel task and our and to identify key components that affect attribution such as doc ID injection strategy and data augmentation. Nonetheless, we believe that our approach and findings are valuable to large-scale pretraining setups. We note that **reviewer Szuz** agrees that “it is easy to see how some version of this method could be incorporated into full-scale LLM training.”\n* **Model size:** Our model is 1.1B and while this is considered small by today’s standards, our computational budget did not allow us to experiment with larger models. We agree that scaling the model size will improve memorization and therefore attribution and we leave such exploration to future work. \n* **Noise in real data:** We hope future research could build on our work and scale this to billion-scale pretraining corpora. Nonetheless, our work is a first step in the direction of building more verifiable LLMs with intrinsic attribution capabilities, and we believe our experiments and findings are valuable to the research community. \n\n* **Results on real data**: To alleviate concerns about the applicability of our approach to real data, we have applied source-aware training to 50K docs from [CCNews](https://huggingface.co/datasets/vblagoje/cc_news). For instruction tuning and evaluation, we prompt ChatGPT to generate question-answer pairs from the articles. We used the exact same training hyperparameters as in the paper. Here are the results: \n\n|                           | In-domain |        | Out-of-domain |        |\n|---------------------------|:---------:|:------:|:-------------:|:------:|\n| Doc ID injection strategy | EM| Hits@1 | EM| Hits@1 |\n| ID-BEGIN| 26.2      | 43.7   | 21.0 | 12.5|\n| ID-END| 24.5 | 46.6   | 23.3| 12.1|\n| ID-REPEAT| 27.8 | 64.7   | 27.8| 36.2|\n\nWe observe the same performance trend as with BioCite regarding the doc ID injection strategies: ID-REPEAT performs best, and our approach achieves reasonable OOD attribution on a pretraining corpus of real data.\n\nWe hope we have addressed your concerns, and we hope you can reconsider your score. Please let us know if you have any more questions!", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "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.\n2. 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).\n3. 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.\n\n2. 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.\n\nGiven 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.", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The authors propose a way to perform attributions to pertaining documents via a model's parametric memory - document IDs are included in the pertaining corpus and then fine-tuning is used to extract these at test time.

The major point of contention for all reviewers is the synthetic nature of the experiments, and whether these experiments demonstrate the viability of the approach in a more 'real world' setting. I concur with the general view expressed by Reviewer Szuz -- the paper presents very interesting ideas, ones that will likely be of interest to others and advance the field, though it is unclear if the evidence presented in this paper demonstrates broader applicability. 

The authors have committed during rebuttals to be more careful and accurate with claims, highlighting the synthetic nature of their setup and the limited conclusions that can be drawn from such a setup. I believe this is sufficient to address the major weakness of the paper - the lack of immediate real world viability is not, in and of itself, a critical factor to reject a paper if the rest of the paper is of significant interest and the claims are accurately made.

[At least one review was discounted during the decision process due to quality]
Decision: Accept