Token Memory Transformer with Infinite Context
摘要
This study proposes an infinite context Transformer model based on Token Memory, which aims to solve the problem of contextual limitation in long text processing. The core of this model is Token Memory, which stores the context memory for each token and provides the information during model generation. The model first splits a long text into segments and then generates local and global attention for each segment. Local attention is generated by the decoder-only Transformer and is discarded after each segment is completed. Global attention is generated by the Token Memory. The Token Memory is retained to provide information for the infinite context, and it is updated after the current segment is calculated. Combining local and global attention can achieve an infinite context. We trained and evaluated the proposed model on the PG-19, C4-en, and BookSum datasets. Our model reaches state-of-the-art results on language modeling tasks with long contexts of up to 1M tokens. Our approach enables large language models to handle infinite contextual tasks, adapting to the increasing demands placed on these models by people.