Source-Aware Training Enables Knowledge Attribution in Language Models

Thank you so much for your feedback. We are glad you found the paper idea clever and innovative! We address your concerns below:

• 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:

We observe that our approach still works reasonably well with random numerical IDs.

• 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.
• 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.
• 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.

i really just am confused as to why this problem was introduced

We assume you refer to the doc ID problem addressed above. If not, please let us know.

• 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.

We hope we have addressed your concerns, and please let us know if there's anything else you’d like us to clarify!

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:

• 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.
• 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.
• 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)
• 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}$]

The 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.

• 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: || OOD Answer EM | OOD Hits@1 | |-|-|-| | Last-names truncated | 88.8%| 47.3%| | Random ID| 95.9%| 39.8%|

We will include results with random IDs in future revisions.

We hope we have addressed your concerns, and please let us know if you have any more questions.

Thank you so much for acknowledging the novelty and importance of our paper! We will address your concerns below:

• 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.”

• 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.

• 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.

• 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. 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:

We 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.

We hope we have addressed your concerns, and we hope you can reconsider your score. Please let us know if you have any more questions!

Thank you for your comments on our paper! We aim to address your comments below:

• 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”.
• 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.
• 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?
• 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.

To alleviate concerns about the applicability of our approach to real data, we have applied source-aware training to 50K docs from CCNews. 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:

We 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.

We hope we have addressed your concerns, and we hope you can reconsider your score. Please let us know of any more concerns!