Retrieval-Augmented Memory for Online Learning
Wenzhang Du
TL;DR
RAM-OL introduces a retrieval-augmented memory mechanism for online learning under concept drift by maintaining a small buffer of past examples and retrieving nearest neighbours in the hidden representation to jointly update on current and neighbour data. The framework includes a naive replay variant and a gated replay variant with time, similarity, and gradient gates to balance fast reuse with robustness to outdated regimes, and it is instantiated on an online MLP. Theoretical framing under a bounded-drift model links retrieval to improved adaptation costs and favorable regret constants in recurring environments; empirically, RAM-OL improves prequential accuracy and greatly reduces seed variance on electricity-related streams, while gated replay remains robust on noisier data. The work highlights practical considerations, including computational overhead, gate design, and privacy implications, and points to future work in scaling, multi-task settings, and secure memory usage.
Abstract
Retrieval-augmented models couple parametric predictors with non-parametric memories, but their use in streaming supervised learning with concept drift is not well understood. We study online classification in non-stationary environments and propose Retrieval-Augmented Memory for Online Learning (RAM-OL), a simple extension of stochastic gradient descent that maintains a small buffer of past examples. At each time step, RAM-OL retrieves a few nearest neighbours of the current input in the hidden representation space and updates the model jointly on the current example and the retrieved neighbours. We compare a naive replay variant with a gated replay variant that constrains neighbours using a time window, similarity thresholds, and gradient reweighting, in order to balance fast reuse of relevant past data against robustness to outdated regimes. From a theoretical perspective, we interpret RAM-OL under a bounded drift model and discuss how retrieval can reduce adaptation cost and improve regret constants when patterns recur over time. Empirically, we instantiate RAM-OL on a simple online multilayer perceptron and evaluate it on three real-world data streams derived from electricity pricing, electricity load, and airline delay data. On strongly and periodically drifting streams, RAM-OL improves prequential accuracy by up to about seven percentage points and greatly reduces variance across random seeds, while on a noisy airline stream the gated variant closely matches the purely online baseline. These results show that retrieval-augmented memory is a practical and robust tool for online learning under concept drift.