New Technologies in Mathematics Seminar

Speaker: Johannes von Oswald (Google)

Title: Retrieval-Centric Deep Learning: Replacing Weight Matrices with Vector Databases

Abstract: The deep learning models powering our current AI revolution are built on parameter matrices that form linear layers. These essential components dictate a learning system’s capacity, as well as its computational and memory requirements during training and inference. In this talk, I will first highlight that linear layers trained by gradient descent can be expressed as key-value memory systems, storing inputs as keys and their corresponding output gradients as values (Irie et al., 2022). This creates a direct connection between gradient-descent-trained linear layers and popular sequence layers that learn in-context, such as Mamba or Gated DeltaNet. Consequently, this duality motivates a gradient-descent-based linear layer counterpart to softmax self-attention. Second, leveraging this duality, I will present both historically overlooked and novel “growing” retrieval-centric layers. These provide a principled drop-in replacement for standard linear layers by storing (roughly) inputs and output gradients in a database, recalling them using non-linear attention. Finally, I will share ongoing empirical results on language modeling that highlight the practical challenges and promising future directions of this new paradigm in retrieval-centric deep learning.