Reproduce “Echo: KV-Cache-Free Associative Recall with Spectral Koopman Operators” on OpenResearch

About this paper

Anupama Sridhar, Alexander Johansen

Echo introduces an architecture that employs Spectral Koopman Attention (SKA) and a Spectral Koopman MLP to achieve associative recall with constant inference memory, addressing the "memory cliff" of State-Space Models and the KV-cache bottleneck of Transformers. The approach achieved 100% retrieval accuracy on the MQAR benchmark and demonstrated competitive language modeling performance, such as surpassing GPT-2 345M on HellaSwag with approximately half the parameters.

Key insights

• Content-addressed retrieval can be explicitly formulated and efficiently solved as a kernel ridge regression problem, where sufficient statistics can be accumulated in constant memory.
• Spectral Koopman operators can effectively identify and retrieve persistent bindings by spectrally normalizing the transition operator and applying a power filter that amplifies persistent eigenmodes.
• The closed-form computation of SKA and norm-preserving operations in the Koopman MLP ensure a full-rank Jacobian, promoting stable and effective gradient flow during training.

Summary by alphaXiv.