Learning to Recall with Transformers Beyond Orthogonal Embeddings

Area of Science:

• Artificial Intelligence
• Machine Learning Theory
• Deep Learning

Background:

• Large language models (LLMs) demonstrate strong knowledge storage and retrieval capabilities, crucial for tasks like question answering.
• Transformers are key to LLM performance, encoding information during training and retrieving it during inference.
• Current theoretical analyses often use idealized assumptions (infinite data, orthogonal embeddings) not reflective of real-world training.

Purpose of the Study:

• To analyze the storage capacity of transformers under realistic training conditions, specifically finite datasets and non-orthogonal embeddings.
• To investigate the impact of gradient descent on a single-layer transformer's ability to learn token-retrieval tasks.
• To derive explicit formulas for storage capacity and understand its scaling properties.

Main Methods:

• Analysis of a single-layer transformer model using empirical gradient descent.
• Focus on the 'early phase' of training on a token-retrieval task with random embeddings.
• Numerical validation of derived theoretical scalings and comparison with statistical lower bounds.

Main Results:

• Derived explicit formulas for the storage capacity of the transformer model.
• Revealed a multiplicative dependence of storage capacity on sample size (N), embedding dimension (d), and sequence length (L).
• Demonstrated that this multiplicative scaling is intrinsic to the problem under non-orthogonal embeddings.

Conclusions:

• The study provides a theoretical understanding of LLM storage capacity in realistic, non-idealized settings.
• The findings highlight the interplay between data, model architecture, and training dynamics in determining knowledge storage.
• The derived scalings offer insights into designing and training more efficient large language models.