Source-linked AI summary
KItCAT: Knowledge Injection via Input Corruption for Auto-regressive Training
Meghanadh Pulivarthi, Kushagra Bhushan, Vineet Kumar, Gaurav Pandey, Jaydeep Sen, Dinesh Raghu, Sachindra Joshi, Yatin Nandwani
TL;DR
Niche documents provide limited repeated exposure to specialized facts, making CPT vulnerable to overfitting and motivating costly paraphrase generation. KItCAT instead corrupts conditioning inputs while preserving targets to create diverse training contexts. It consistently outperforms standard CPT across four benchmarks and three model families, while remaining complementary to paraphrase augmentation.
Problem
Niche documents seldom repeat facts, limiting CPT’s repeated exposure to new knowledge and making robust acquisition difficult; paraphrasing can be costly and require powerful external LLMs.
Method
KItCAT stochastically corrupts input sequences and trains the model to predict unchanged targets under perturbed contexts, generating diversity without external synthetic data.
Results
KItCAT consistently outperforms standard next-token prediction-based training across four knowledge-injection benchmarks and three model families.
Takeaways & Limitations
KItCAT is lightweight, architecture-agnostic, complementary to existing augmentation, and incurs negligible computational overhead without external paraphrasing models.
Takeaways & Limitations
KItCAT reduces but does not fully replace synthetic paraphrase generation, and the highest downstream performance comes from combining both augmentations.
Abstract
from arXiv · showhide
LLMs acquire vast amounts of knowledge during pre-training, but often lack the specialized knowledge needed to answer questions from niche sources such as manuals or technical documents unseen during pre-training. Continued pre-training (CPT) is widely used to inject such knowledge into model parameters. However, niche documents seldom repeat facts, making it difficult for CPT to robustly acquire such knowledge. Recent works address this by generating multiple paraphrases of the new knowledge, but paraphrasing is computationally expensive and typically requires powerful LLMs. In this work, we introduce KItCAT: Knowledge Injection via Corrupted Auto-regressive Training, a lightweight training strategy that reduces the need for paraphrasing in decoder-only LLMs. KItCAT augments standard next-token prediction by stochastically corrupting the input sequence. During training, a random subset of input tokens is replaced with other vocabulary tokens while the original next-token labels are kept unchanged. This simple intervention generates diverse training inputs from each sample, enabling large-scale data augmentation at negligible cost. We show that KItCAT consistently improves over CPT across multiple datasets and model families. Code is available at https://github.com/meghanadhpulivarthi/KItCAT.
1 Introduction
LLMs often miss specialized knowledge in niche documents, where low repetition limits CPT and can encourage surface-pattern overfitting. KItCAT addresses this by corrupting conditioning inputs to create diverse training contexts while preserving next-token labels.
- Specialized knowledge absent or underrepresented on the public web is often missing from pretrained LLMs, motivating CPT for proprietary manuals and technical documents.
- Niche documents seldom repeat facts, limiting repeated exposure during CPT and making robust internalization difficult while encouraging lexical-pattern overfitting.
- Paraphrase-based augmentation improves diversity but can be costly, slow, dependent on large proprietary LLMs, and infeasible under privacy restrictions.
- KItCAT prevents identical inputs across epochs by applying controlled corruption to input sequences, generating diverse contexts without external synthetic data generation.
- The four corruption schemes modify different input positions while retaining standard next-token training labels.
- Across multiple benchmarks and three model families, KItCAT consistently improves CPT, reduces reliance on paraphrasing, and remains complementary to it.
2 Related Work
Knowledge injection methods include CPT and SFT approaches, and both benefit from diverse training data. Corruption can target inputs or hidden representations across these training styles.
- CPT injects knowledge from raw domain text and paraphrases, whereas SFT transforms raw text into task-oriented formats such as QA, summarization, or reading comprehension.
- Both CPT- and SFT-based approaches benefit from high-diversity training data.
- Corruption methods can modify input tokens or hidden representations and can be used in both SFT and CPT.
- Dropout corrupts hidden activations, while Latent Paraphrasing perturbs hidden representations for knowledge injection.
3 The KItCAT approach
KItCAT trains decoder-only LLMs on corrupted prefixes while keeping target tokens unchanged, encouraging predictions that remain stable across perturbed contexts. Four stochastic schemes provide different corruption patterns.
- KItCAT samples corrupted prefixes from Q(˜s | s) while preserving the target token as a likely continuation.
- Training on perturbed prefixes encourages the model to use information that remains present across perturbations rather than spurious lexical patterns.
- KItCAT-mask uses a special mask token, KItCAT-rand uses random vocabulary tokens, KItCAT-SSMBA uses plausible alternatives, and KItCAT-MASKER masks informative keywords.
- For each corruption method, selected token positions are replaced by r_j while unselected positions retain s_j.
- KItCAT-mask, KItCAT-rand, and KItCAT-SSMBA independently select token positions with probability p and resample positions and replacements each epoch.
- KItCAT-mask, KItCAT-rand, KItCAT-SSMBA, and KItCAT-MASKER use mask tokens, uniform vocabulary samples, a frozen masked language model, and TF–IDF keyword spans, respectively.
- The four variants replace conditioning-prefix tokens only, leaving the target token s_t unchanged.
4 Experimental Setup
The study injects knowledge from four corpora, evaluates answers with an LLM judge, and compares standard NTP with four KItCAT variants across multiple model families.
- Knowledge is injected from a PopQA subset, Companies, and two Redbooks, with evaluation using each dataset’s accompanying test question-answer pairs.
- The comparison includes the out-of-the-box Instruct model, standard CPT using next-token prediction, and four KItCAT corruption schemes.
- KItCAT-SSMBA uses an encoder model for contextually plausible replacements, while KItCAT-MASKER preferentially masks informative keywords.
- The main experiments use LoRA fine-tuning of Mistral-7B-Instruct-v0.3, with Qwen3-14B and Llama-2-7B-Chat testing model-family and scale generalization.
5 Experimental Results
Across knowledge-injection experiments, KItCAT improves over standard NTP, remains effective across model families, and complements synthetic data augmentation while offering substantial compute savings. Its training dynamics also indicate greater stability than NTP.
- All four KItCAT variants outperform NTP on average across the four knowledge-injection corpora.
- 27.7 is KItCAT-rand’s average score, improving 10.0 points over NTP’s 17.6; KItCAT-mask reaches 26.1.
- KItCAT generalizes to Qwen3-14B and Llama-2-7B-Chat, consistently improving over NTP across the evaluated corpora.
- KItCAT is complementary to synthetic data augmentation, and original-data KItCAT achieves higher scores than NTP with substantially more Rephrase or Entigraph data.Using only original data, KItCAT scores 26.84 versus 22.35 for NTP with 9× Rephrase data and 24.53 with 5× Entigraph data, a 5–9× effective data multiplier.
- KItCAT-mask matches Rephrase-10× accuracy at 15.2% of its estimated compute cost on the Companies dataset.KItCAT-mask achieves 27.6 accuracy at 30 PFLOPs, compared with Rephrase-10× at 27.4 accuracy and 198 PFLOPs.
- NTP overfits after reaching minimum validation loss within seven epochs, whereas KItCAT-mask and KItCAT-rand maintain stable validation loss throughout training.The validation-loss comparison uses QA pairs from the Companies validation set.
6 Conclusion
The conclusion presents KItCAT as a lightweight modification to autoregressive training that diversifies conditioning inputs for knowledge injection. Across four benchmarks and three model families, it outperforms standard NTP while avoiding external paraphrasing models and negligible computational overhead.
- KItCAT modifies autoregressive training so the model does not encounter the exact same input twice.The stated purpose is to mitigate reinforcement of spurious lexical correlations in low-diversity CPT settings.
- Across four benchmarks and three model families, KItCAT consistently outperforms standard next-token prediction-based training.
- KItCAT is described as lightweight, architecture-agnostic, complementary to existing augmentation, and applicable without external paraphrasing models.The conclusion also characterizes its computational overhead as negligible.
Limitations
KItCAT reduces reliance on synthetic paraphrase generation but does not fully replace it; combining both augmentation methods achieves the highest downstream performance.
- KItCAT does not fully substitute for expensive synthetic paraphrase generation.Synthesized paraphrases continue to provide complementary benefits, and the highest downstream performance comes from combining corruption with paraphrase-based augmentation.
A.1 Example Model Responses
The example response answers a factual question about DoodleDreamDelight Publishing Ltd. by identifying its founder and previous profession.
- Emily Chen founded DoodleDreamDelight Publishing Ltd. and was previously a kindergarten teacher.
NTP
The baseline prediction is incorrect: it attributes DoodleDreamDelight Publishing Ltd. to Emily Carter, a former graphic designer.
- The baseline answer incorrectly names Emily Carter, previously a graphic designer in advertising, as the founder.
KItCAT-mask
KItCAT augments decoder-only autoregressive training with label-preserving input corruption, constraining predictions to remain stable under perturbed prefixes. Experiments report gains across model families, with modest added training cost and an intermediate corruption range performing best.
- Motivation: The objective addresses overfitting to spurious prefix-target correlations, especially when low-data training sequences are repeatedly observed.The motivation is to stabilize the predictive distribution under corrupted input contexts.
- Method: The corruption constraint trades off fitting training data against prediction stability through a nonnegative Lagrange multiplier λ.λ = 0 recovers standard autoregressive training, while different λ choices yield related objectives.
- Method: KItCAT constrains next-token prediction to vary smoothly under label-preserving perturbations of the input prefix.The formulation bounds expected deviation between predictions on original and corrupted contexts.
- Results: Gains extend beyond one model family, while reconstruction-based and keyword-selection schemes do not consistently outperform simple random corruption.Additional Qwen3-14B and Llama-2-7B-Chat experiments use the original 1× corpus.
- Efficiency: Input corruption skips paraphrase generation and adds only a modest increase in training cost.The compute estimates separate one-time paraphrase-generation FLOPs from LoRA training FLOPs.
- Results: Performance peaks around p = 0.15–0.30 and declines when corruption is too weak or too strong.The sensitivity is more pronounced on the data-scarce Companies corpus, where excessive corruption can confuse learning.