EnergyGradMem: Learning an Energy Function for Test-Time Memory Writing

Compressing context into a compact memory state lets a language model answer future queries without retaining or reprocessing it. GradMem writes such memories by test-time optimization of a context-reconstruction objective chosen a priori. We introduce EnergyGradMem, which instead learns a query-independent energy function from downstream supervision. At inference, it writes to memory by minimizing the energy over a context, and uses only memory and query for inference. In multi-query associative recall (MQAR) experiments we compare learned energy with reconstruction-based memory updates, refinement and longer-context generalization.