Memory Layers - Meta Open Source: Methods for Increasing Parameters of Large Models Without Increasing Computing Power
Memory Layers, introduced by Meta, are a method for adding extra parameters to a model using a trainable key-value lookup mechanism without increasing the number of floating-point operations (FLOPs). They supplement computationally intensive feedforward layers with sparse activation, providing...
What are Memory Layers?
Memory Layers, introduced by Meta, adds extra parameters to a model using a trainable key-value lookup mechanism without increasing floating-point operations (FLOPs). Based on sparse activation, Memory Layers supplement computationally intensive feedforward layers, providing dedicated capacity for inexpensive storage and retrieval of information. Memory Layers have shown significant practicality in large-scale language models, especially in handling factual tasks, significantly improving model performance. By replacing one or more feedforward networks (FFNs) with memory layers in a Transformer layer, the model's factual accuracy and knowledge acquisition capabilities are significantly improved without significantly increasing computational costs.
Main functions of Memory Layers
- Parameter expansionIt supports adding additional parameters to the model without increasing the computational burden (i.e., without increasing FLOPs), thus expanding the model's capacity.
- Information storage and retrievalProvides a dedicated mechanism for storing and retrieving information.
- Supplement dense layerComplementing computationally intensive feed-forward layers, it provides dedicated capacity for information storage and retrieval based on sparse activation patterns.
- Improve the accuracy of factsIn factual tasks, it significantly improves the accuracy of language models, enabling them to more accurately remember and use facts from the training data.
- Performance improvement: To enable the model to outperform traditional dense models and hybrid expert models under conditions of computation and parameter matching.
Technical principles of Memory Layers
- Key-value lookup mechanismIt operates based on a simple key-value lookup mechanism, where both keys and values are encoded as embedding vectors. Given a query vector q, and a set of keys K and values V, the output is a soft combination of values, with weights determined by the similarity between q and the corresponding keys.
- Sparse activationUnlike traditional attention mechanisms, keys and values in Memory Layers are trainable parameters, not activation values. Since the number of keys and values is typically large, Memory Layers require sparse lookups and updates, with only the top-k keys and their corresponding values most similar to the query participating in the output computation.
- Product Key SearchTo address the bottleneck of query-key retrieval in large-scale memory, Memory Layers employs product-quantized keys, which efficiently perform top-k lookups of the entire key set using two smaller sets of keys.
- Parallelization implementationBecause Memory Layers have high memory requirements, large-scale scaling is achieved by parallelizing embedded lookup and aggregation across multiple GPUs.
- Shared memory parameters: Using a shared memory parameter pool, performance can be increased based on multiple memory layers without increasing the number of parameters.
Project address for Memory Layers
- Project official website:https://ai.meta.com/research/publications/memory-layers-at-scale/
- GitHub repository:https://github.com/facebookresearch/memory
- arXiv technical paper:https://arxiv.org/pdf/2412.09764
Application scenarios of Memory Layers
- Question and Answer SystemIn question-answering systems, this helps the model remember and retrieve large amounts of factual information to provide accurate answers.
- Language ModelIn natural language processing, enhancing the memory capacity of language models makes them more effective at handling long texts and complex dialogues.
- Recommendation systemThe purpose of a recommendation system is to store users’ historical behavior and preferences to provide more personalized recommendations.
- knowledge graphIn knowledge graph applications, relationships between entities are stored and retrieved, supporting complex queries and reasoning.
- Dialogue systemIn dialogue systems, this helps the model remember the dialogue history, making the dialogue more coherent and natural.