Elamur Research⌕
☰
home / papers / 2025 / Towards Context-Robust LLMs: A Gated Representation Fine-tuning Approach

Towards Context-Robust LLMs: A Gated Representation Fine-tuning Approach

Paper · 2025 · Shenglai Zeng, Pengfei He, Kai Guo et al.

Summary

LLMs that receive retrieved passages, as in RAG, tend to trust that text too much. Misleading or off-topic passages can override knowledge the model already has. The authors define a "context-robust" LLM by four behaviours:

  • use the context when the model lacks the answer;
  • accept the context when it agrees with the model's knowledge;
  • report both answers when the context conflicts with internal knowledge;
  • ignore passages that do not help.

Their method, Grft, is a plug-in that edits hidden-layer representations. It has two parts. A sigmoid gate estimates whether an input is abnormal, meaning contradictory or unhelpful. Low-rank adapters then shift the representation inside a linear subspace, scaled by the gate value. An optional variant, Grft-requery, asks the model again when the output signals a conflict or useless context.

Grft is trained on under 200 questions, mostly from ConflictQA, and tested on Llama-2-7B and Llama-3-8B-instruct. It improves accuracy on misleading and unhelpful contexts while keeping performance on helpful ones. This makes it a cheap way to harden RAG pipelines.

Key points

  • Parameter count: Grft trains about 36.9K parameters, roughly 0.0004–0.0005% of a 7B model. By comparison, LoRA (r=4) uses 2.1M and full fine-tuning uses 6.74B. The gate adds only about 4.1K parameters on top of ReFT.
  • Gate behaviour: average gate values are above 0.7 for contradictory contexts and close to 1 for unhelpful contexts with known questions. For normal inputs (matching contexts, or unknown questions with helpful contexts), they stay below 0.3.
  • Transfer to COUNTERFACT (1,000 samples, Llama-7B-Chat): with contradictory contexts, Grft reaches 62.52% versus 42.24% for the plain LLM. Grft-requery reaches 75.78%, which is 24.35 points above Astute-RAG.
  • Transfer to a 1,200-question NQ subset: Grft gains 9.71 points on distracted contexts and 24.11 points on random contexts. Grft-requery reaches 83.13% and 82.20% on these two context types.
  • Ablations: removing the gate or the gate loss hurts accuracy on helpful contexts. The full Grft stays close to the original LLM on those inputs.
  • Limitations: the work only covers single context–question pairs and two Llama models.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢AmazonSeattle · United States112 works · 40,565 citations🎓Michigan State UniversityEast Lansing · United States23 works · 1,077 citations

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 worksInformation Systems · 1,437 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Shenglai Zeng, Pengfei He, Kai Guo, Tianqi Zheng, Hanqing Lu, Yue Xing and Hui Liu. 2025. Towards Context-Robust LLMs: A Gated Representation Fine-tuning Approach. Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2025.acl-long.506

@inproceedings{zeng2025context,
  title = {{Towards Context-Robust LLMs: A Gated Representation Fine-tuning Approach}},
  author = {Shenglai Zeng, Pengfei He, Kai Guo, Tianqi Zheng, Hanqing Lu, Yue Xing and Hui Liu},
  booktitle = {Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  year = {2025},
  doi = {10.18653/v1/2025.acl-long.506},
  url = {https://doi.org/10.18653/v1/2025.acl-long.506}
}

References

34 references, 17 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Lameck Mbangula Amugongo, Pietro Mascheroni, Steven Geoffrey Brooks, Stefan Doering, and Jan Seidel. 2024. Retrieval augmented generation for large language models in healthcare: A systematic review. 2024a. The power of noise: Redefining retrieval for rag systems. In Proceedings of the 47th Internat…
  2. Florin Cuconasu, Giovanni Trappolini, Federico Siciliano, Simone Filice, Cesare Campagnano, Yoelle Maarek, Nicola Tonellotto, and Fabrizio Silvestri. 2024b. The power of noise: Redefining retrieval for rag systems. In Proceedings of the 47th International ACM SIGIR Conference on Research and Develop…
  3. Gelei Deng, Yi Liu, Kailong Wang, Yuekang Li, Tianwei Zhang, and Yang Liu. 2024. Pandora: Jailbreak gpts by retrieval augmented generation poisoning. NDSS Workshop on AI Systems with Confidential Computing (AISCC). find
  4. Wenqi Fan, Yujuan Ding, Liangbo Ning, Shijie Wang, Hengyun Li, Dawei Yin, Tat-Seng Chua, and Qing Li. 2024. A survey on rag meeting llms: Towards retrieval-augmented large language models. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 6491– 6501.
  5. Feiteng Fang, Yuelin Bai, Shiwen Ni, Min Yang, Xiaojun Chen, and Ruifeng Xu. 2024. Enhancing noise robustness of retrieval-augmented language models with adaptive adversarial training. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers),…
  6. Yunfan Gao, Yun Xiong, Xinyu Gao, Kangxiang Jia, Jinliu Pan, Yuxi Bi, Yi Dai, Jiawei Sun, and Haofen Wang. 2023. Retrieval-augmented generation for large language models: A survey. arXiv preprint arXiv:2312.10997.
  7. Debra L Hollister, Avelino Gonzalez, and James Hollister. 2017. Contextual reasoning in human cognition and the implications for artificial intelligence systems. In Modeling and Using Context: 10th International and Interdisciplinary Conference, CONTEXT 2017, Paris, France, June 20-23, 2017, Proceed… find
  8. Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-rag: Learning to retrieve, generate, and critique through self-reflection. In The Twelfth International Conference on Learning Representations.
  9. 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. 2020. Retrieval-augmented generation for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33:9459–9474.
  10. Jiawei Chen, Hongyu Lin, Xianpei Han, and Le Sun. 2024. Benchmarking large language models in retrieval-augmented generation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 17754–17762.
  11. Yuping Lin, Pengfei He, Han Xu, Yue Xing, Makoto Yamada, Hui Liu, and Jiliang Tang. 2024. Towards understanding jailbreak attacks in llms: A representation space analysis. arXiv preprint arXiv:2406.10794. find
  12. Florin Cuconasu, Giovanni Trappolini, Federico Siciliano, Simone Filice, Cesare Campagnano, Yoelle Maarek, Nicola Tonellotto, and Fabrizio Silvestri. find
  13. Yanming Liu, Xinyue Peng, Xuhong Zhang, Weihao Liu, Jianwei Yin, Jiannan Cao, and Tianyu Du. 2024. Ra-isf: Learning to answer and understand from retrieval augmentation via iterative self-feedback. arXiv preprint arXiv:2403.06840. find
  14. Alex Mallen, Akari Asai, Victor Zhong, Rajarshi Das, Daniel Khashabi, and Hannaneh Hajishirzi. 2023. When not to trust language models: Investigating effectiveness of parametric and non-parametric memories. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (V… find
  15. Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. 2022. Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems, 35. find
  16. Shiyu Ni, Keping Bi, Jiafeng Guo, and Xueqi Cheng. 2024. When do LLMs need retrieval augmentation? mitigating LLMs’ overconfidence helps retrieval augmentation. In Findings of the Association for Computational Linguistics ACL 2024, pages 11375–11388, Bangkok, Thailand and virtual meeting. Associatio… find
  17. Ruiyang Ren, Yuhao Wang, Yingqi Qu, Wayne Xin Zhao, Jing Liu, Hao Tian, Hua Wu, Ji-Rong Wen, and Haifeng Wang. 2023. Investigating the factual knowledge boundary of large language models with retrieval augmentation. arXiv preprint arXiv:2307.11019.
  18. Kunal Sawarkar, Abhilasha Mangal, and Shivam Raj Solanki. 2024. Blended rag: Improving rag (retriever-augmented generation) accuracy with semantic search and hybrid query-based retrievers. arXiv preprint arXiv:2404.07220. — [arXiv:2404.07220](https://arxiv.org/abs/2404.07220) find
  19. Spurthi Setty, Harsh Thakkar, Alyssa Lee, Eden Chung, and Natan Vidra. 2024. Improving retrieval for rag based question answering models on financial documents. arXiv preprint arXiv:2404.07221. — [arXiv:2404.07221](https://arxiv.org/abs/2404.07221) find
  20. Fei Wang, Xingchen Wan, Ruoxi Sun, Jiefeng Chen, and Sercan Ö Arık. 2024. Astute rag: Overcoming imperfect retrieval augmentation and knowledge conflicts for large language models. arXiv preprint arXiv:2410.07176.
  21. Yike Wang, Shangbin Feng, Heng Wang, Weijia Shi, Vidhisha Balachandran, Tianxing He, and Yulia Tsvetkov. 2023a. Resolving knowledge conflicts in large language models. arXiv preprint arXiv:2310.00935.
  22. Nirmalie Wiratunga, Ramitha Abeyratne, Lasal Jayawardena, Kyle Martin, Stewart Massie, Ikechukwu Nkisi-Orji, Ruvan Weerasinghe, Anne Liret, and Bruno Fleisch. 2024. Cbr-rag: case-based reasoning for retrieval augmented generation in llms for legal question answering. In International Conference on C…
  23. Zhengxuan Wu, Aryaman Arora, Zheng Wang, Atticus Geiger, Dan Jurafsky, Christopher D. Manning, and Christopher Potts. 2024a. ReFT: Representation finetuning for language models. find
  24. Zhengxuan Wu, Aryaman Arora, Zheng Wang, Atticus Geiger, Dan Jurafsky, Christopher D Manning, and Christopher Potts. 2024b. Reft: Representation finetuning for language models. arXiv preprint arXiv:2404.03592. find
  25. Jian Xie, Kai Zhang, Jiangjie Chen, Renze Lou, and Yu Su. Adaptive chameleon or stubborn sloth: Revealing the behavior of large language models in knowledge conflicts. In The Twelfth International Conference on Learning Representations. Ori Yoran, Tomer Wolfson, Ori Ram, and Jonathan Berant. Making…
  26. Shenglai Zeng, Jiankun Zhang, Pengfei He, Yue Xing, Yiding Liu, Han Xu, Jie Ren, Shuaiqiang Wang, Dawei Yin, Yi Chang, et al. 2024a. The good and the bad: Exploring privacy issues in retrieval-augmented generation (rag). ACL Findings. find
  27. Shenglai Zeng, Jiankun Zhang, Bingheng Li, Yuping Lin, Tianqi Zheng, Dante Everaert, Hanqing Lu, Hui Liu, Yue Xing, Monica Xiao Cheng, et al. 2024b. Towards knowledge checking in retrieval-augmented generation: A representation perspective. arXiv preprint arXiv:2411.14572.
  28. Siyun Zhao, Yuqing Yang, Zilong Wang, Zhiyuan He, Luna K Qiu, and Lili Qiu. 2024. Retrieval augmented generation (rag) and beyond: A comprehensive survey on how to make your llms use external data more wisely. arXiv preprint arXiv:2409.14924. — [arXiv:2409.14924](https://arxiv.org/abs/2409.14924) find
  29. Chujie Zheng, Fan Yin, Hao Zhou, Fandong Meng, Jie Zhou, Kai-Wei Chang, Minlie Huang, and Nanyun Peng. 2024. Prompt-driven llm safeguarding via directed representation optimization. arXiv preprint arXiv:2401.18018. find
  30. Yile Wang, Peng Li, Maosong Sun, and Yang Liu. 2023b. Self-knowledge guided retrieval augmentation for large language models. In Findings of the Association for Computational Linguistics: EMNLP 2023, pages 10303–10315.
  31. Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo, Richard Ren, Alexander Pan, Xuwang Yin, Mantas Mazeika, Ann-Kathrin Dombrowski, et al. 2023. Representation engineering: A top-down approach to ai transparency. arXiv preprint arXiv:2310.01405.
  32. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837.
  33. Wei Zou, Runpeng Geng, Binghui Wang, and Jinyuan Jia. 2024. Poisonedrag: Knowledge poisoning attacks to retrieval-augmented generation of large language models. find
  34. USENIX Security 2025. *Figure 5: Ablation study on i-th layer intervention* *Figure 6: Ablation study on number of training queries* In this section, we conduct ablation studies on the intervention effect on different layers and the minimal requirement of training samples. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/zeng-2025-context-robust-llms-gated-representation/