Language Model Memory and Memory Models for Language
Abstract
The ability of machine learning models to store input information in hidden layer embeddings, a form of model `memory', is widely employed but not well characterized. We find that causal language model embeddings typically contain relatively little input information regardless of data and compute scales 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 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 conclude that next token prediction training alone is poorly suited for accurate memory formation, motivating the use of combined objective functions for models where the entire input is not exposed.