FastMem: Fast Memorization of Prompt Improves Context Awareness of Large Language Models

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Zhu, Junyi, Liu, Shuochen, Yu, Yu, Tang, Bo, Yan, Yibo, Li, Zhiyu, Xiong, Feiyu, Xu, Tong, Blaschko, Matthew B.
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917794926297088
author Zhu, Junyi
Liu, Shuochen
Yu, Yu
Tang, Bo
Yan, Yibo
Li, Zhiyu
Xiong, Feiyu
Xu, Tong
Blaschko, Matthew B.
author_facet Zhu, Junyi
Liu, Shuochen
Yu, Yu
Tang, Bo
Yan, Yibo
Li, Zhiyu
Xiong, Feiyu
Xu, Tong
Blaschko, Matthew B.
contents Large language models (LLMs) excel in generating coherent text, but they often struggle with context awareness, leading to inaccuracies in tasks requiring faithful adherence to provided information. We introduce FastMem, a novel method designed to enhance instruction fine-tuned LLMs' context awareness through fast memorization of the prompt. FastMem maximizes the likelihood of the prompt before inference by updating only the last Feed-Forward Network (FFN) module. This targeted approach ensures efficient optimization without overfitting, significantly improving the model's ability to comprehend and accurately follow the context. Our experiments demonstrate substantial gains in reading comprehension, text summarization and adherence to output structures. For instance, FastMem improves the accuracy of Llama 3-8B-Inst on the NQ-SWAP dataset from 59.1% to 71.6%, and reduces the output structure failure rate of Qwen 1.5-4B-Chat from 34.9% to 25.5%. Extensive experimental results highlight FastMem's potential to offer a robust solution to enhance the reliability and accuracy of LLMs in various applications. Our code is available at: https://github.com/IAAR-Shanghai/FastMem
format Preprint
id arxiv_https___arxiv_org_abs_2406_16069
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle FastMem: Fast Memorization of Prompt Improves Context Awareness of Large Language Models
Zhu, Junyi
Liu, Shuochen
Yu, Yu
Tang, Bo
Yan, Yibo
Li, Zhiyu
Xiong, Feiyu
Xu, Tong
Blaschko, Matthew B.
Computation and Language
Artificial Intelligence
Large language models (LLMs) excel in generating coherent text, but they often struggle with context awareness, leading to inaccuracies in tasks requiring faithful adherence to provided information. We introduce FastMem, a novel method designed to enhance instruction fine-tuned LLMs' context awareness through fast memorization of the prompt. FastMem maximizes the likelihood of the prompt before inference by updating only the last Feed-Forward Network (FFN) module. This targeted approach ensures efficient optimization without overfitting, significantly improving the model's ability to comprehend and accurately follow the context. Our experiments demonstrate substantial gains in reading comprehension, text summarization and adherence to output structures. For instance, FastMem improves the accuracy of Llama 3-8B-Inst on the NQ-SWAP dataset from 59.1% to 71.6%, and reduces the output structure failure rate of Qwen 1.5-4B-Chat from 34.9% to 25.5%. Extensive experimental results highlight FastMem's potential to offer a robust solution to enhance the reliability and accuracy of LLMs in various applications. Our code is available at: https://github.com/IAAR-Shanghai/FastMem
title FastMem: Fast Memorization of Prompt Improves Context Awareness of Large Language Models
topic Computation and Language
Artificial Intelligence
url https://arxiv.org/abs/2406.16069