Language Model Memory and Memory Models for Language
This addresses the inefficiency in memory formation for language models, which is incremental but could benefit applications requiring accurate input storage and retrieval.
The paper tackles the problem of language models storing input information in embeddings, finding that standard language model embeddings contain little input information regardless of training scale, while autoencoder embeddings achieve nearly perfect memory. It introduces a parallelizable encoder-decoder architecture with combined objective functions to improve memory formation, leading to substantial computational efficiencies.
The ability of machine learning models to store input information in hidden layer vector embeddings, analogous to the concept of `memory', is widely employed but not well characterized. We find that language model embeddings typically contain relatively little input information regardless of data and compute scale during training. In contrast, embeddings from autoencoders trained for input regeneration are capable of nearly perfect memory formation. The substitution of memory embeddings for token sequences leads to substantial computational efficiencies, motivating the introduction of a parallelizable encoder-decoder memory model architecture. Upon causal training these models contain information-poor embeddings incapable of arbitrary information access, but by combining causal and information retention objective functions they learn to form and decode information-rich memories. Training can be further streamlined by freezing a high fidelity encoder followed by a curriculum training approach where decoders first learn to process memories and then learn to additionally predict next tokens. We introduce the perspective that next token prediction training alone is poorly suited for accurate memory formation as the objective itself is non-invertible, motivating the use of combined objective functions for models where the entire input is not exposed.