Explicit Memory Learning with Expectation Maximization

Large Language Models (LLMs) have revolutionized the landscape of natural language processing, demonstrating remarkable abilities across various complex tasks. However, their stateless nature limits the capability to retain information across interactions, hindering performance in scenarios requiring historical context recall. To mitigate this, current approaches primarily use explicit memory to allow LLMs to store useful information, which is accessible, readable, and interpretable. Nevertheless, explicit memory lacks the reliable learning mechanisms of implicit memory, which can be optimized end-to-end. To harness the benefits of both, we introduce EM 2 , a novel framework enhancing explicit memory updates via the Expectation-Maximization (EM) algorithm. EM 2 treats memory as a latent variable, ensuring continual learning and improvement during updates. Experimental results on streaming inference tasks demonstrate that EM 2 outperforms existing methods without memory or with static external memory. Our in-depth analysis highlights that EM 2 significantly enhances performance across various backbones and memory strategies, providing a robust solution for advancing LLM memory management and enabling explicit memory to learn and improve similarly to implicit memory.