# Embedding Domain Knowledge for Large Language Models via Reinforcement Learning from Augmented Generation

arXiv: https://arxiv.org/abs/2509.20162 · HTML: https://mlsanity.com/paper/2509.20162/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 43 references that are in the map.
- First posted on arXiv 24 Sep 2025.

## Builds on (most cited first, 24 of 24)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Finetuned Language Models Are Zero-Shot Learners](https://mlsanity.com/paper/2109.01652/) (2021)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [REALM: Retrieval-Augmented Language Model Pre-Training](https://mlsanity.com/paper/2002.08909/) (2020)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Exams](https://mlsanity.com/paper/2009.13081/) (2020)
- [SimPO: Simple Preference Optimization with a Reference-Free Reward](https://mlsanity.com/paper/2405.14734/) (2024)
- [A Multitask, Multilingual, Multimodal Evaluation of ChatGPT on Reasoning, Hallucination, and Interactivity](https://mlsanity.com/paper/2302.04023/) (2023)
- [How Much Knowledge Can You Pack Into the Parameters of a Language Model?](https://mlsanity.com/paper/2002.08910/) (2020)
- [NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models](https://mlsanity.com/paper/2405.17428/) (2024)
- [Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning](https://mlsanity.com/paper/2305.14160/) (2023)
- [KnowPrompt: Knowledge-aware Prompt-tuning with Synergistic Optimization for Relation Extraction](https://mlsanity.com/paper/2104.07650/) (2021)
- [Fine-Tuning or Retrieval? Comparing Knowledge Injection in LLMs](https://mlsanity.com/paper/2312.05934/) (2023)
- [SAC3: Reliable Hallucination Detection in Black-Box Language Models via Semantic-aware Cross-check Consistency](https://mlsanity.com/paper/2311.01740/) (2023)
- [Crawling the Internal Knowledge-Base of Language Models](https://mlsanity.com/paper/2301.12810/) (2023)
- [Injecting Domain-Specific Knowledge into Large Language Models: A Comprehensive Survey](https://mlsanity.com/paper/2502.10708/) (2025)
- [Large Language Model-Aware In-Context Learning for Code Generation](https://mlsanity.com/paper/2310.09748/) (2023)
- [A Reasoning-Focused Legal Retrieval Benchmark](https://mlsanity.com/paper/2505.03970/) (2025)
- [m2mKD: Module-to-Module Knowledge Distillation for Modular Transformers](https://mlsanity.com/paper/2402.16918/) (2024)

## Built on by (most cited first, 1 of 1)

- [GRASP: GRanularity-Aware Search Policy for Agentic RAG](https://mlsanity.com/paper/2607.10463/) (2026)

## Abstract (arXiv)

Large language models (LLMs) often exhibit limited performance on domain-specific tasks due to the natural disproportionate representation of specialized information in their training data and the static nature of these datasets. Knowledge scarcity and temporal lag create knowledge gaps for domain applications. While post-training on domain datasets can embed knowledge into models, existing approaches have some limitations. Continual Pre-Training (CPT) treats all tokens in domain documents with equal importance, failing to prioritize critical knowledge points, while supervised fine-tuning (SFT) with question-answer pairs struggles to develop the coherent knowledge structures necessary for complex reasoning tasks. To address these challenges, we propose Reinforcement Learning from Augmented Generation (RLAG). Our approach iteratively cycles between sampling generations and optimizing the model through calculated rewards, effectively embedding critical and contextually coherent domain knowledge. We select generated outputs with the highest log probabilities as the sampling result, then compute three tailored reward metrics to guide the optimization process. To comprehensively evaluate domain expertise, we assess answer accuracy and the rationality of explanations generated for correctly answered questions. Experimental results across medical, legal, astronomy, and current events datasets demonstrate that our proposed method significantly outperforms baseline approaches. Our code and data are open sourced at https://github.com/ChaojunNie/RLAG.