Reservoir Sampling

A customer-service model learns from messages that arrive one after another. You want to keep a few of them for later replay, but you have only four slots in memory. You also do not know how many more messages will come: twenty or two million.

Keep the first four? Then the memory knows only the beginning. Always hold the last four? Then it knows only the end. A fair draw of “four out of the whole” would require storing everything or knowing the number of messages in advance. The question is: with a fixed number of slots and a stream of unknown length, how do you give every message the same chance?

Reservoir Sampling is a way of maintaining a random sample of fixed size during a single pass over data whose count we do not know in advance. After each new message, the memory, called a reservoir here, contains a fair random sample of everything that has arrived so far. The simplest version, Algorithm R, is described by Vitter, Random Sampling with a Reservoir, §2, who attributes it to Alan Waterman.

The recipe in two steps

1. The first four messages go into memory without any draw.
2. Message number $t$ (the fifth, the sixth and every later one) gets into memory with chance $4/t$. If it gets in, it replaces one of the four stored messages, chosen at random.

So the fifth message gets in with chance 4/5, the hundredth with chance 4/100. New messages have a smaller and smaller chance of getting in, but the old ones have already been exposed to being thrown out many times, and these two effects balance exactly. Let us check this after the fifth message. It gets in with chance 4/5. The first message drops out only if the fifth got in (4/5) and landed on exactly its slot (1/4), that is, with chance 1/5. It stays with chance 4/5: the same as the fifth has. After $t$ messages, each of them is in memory with chance $4/t$.

An equal chance is not a guarantee

In the experiment, the stream has 20 messages. Eighteen are about delivery, and two, the ninth and the fifteenth, are about returns. Each run sends such a stream through 2000 times and counts how often a given position stayed in memory at the end. The draws are real, so the results differ by a fraction of a percent each time.

With Reservoir Sampling, every bar is at about 20%, that is, 4/20. At the same time, in most runs neither of the two returns messages stays in memory. The exact value is 3060/4845, about 63%: that is how many of all the foursomes chosen from 20 messages consist solely of the eighteen delivery messages. The method is fair to individual messages, but it does not make sure that every category has a representative.

Where this matters

In training with replay, Reservoir Sampling is sometimes the way of writing to the buffer: you do not need to know the length of the stream or the boundaries between tasks. Chaudhry et al., Continual Learning with Tiny Episodic Memories, §2, Algorithm 2 use it in exactly this way. In their experiments on image tasks it worked well, except with very small memories, where earlier classes disappeared from the buffer; strategies that ensure equal representation of classes did better then (§4.4). This is a result for image classification, not for LLMs.

The algorithm itself is not about machine learning. Vitter's 1985 paper describes sampling records from a magnetic tape of unknown length and proposes faster variants which, instead of drawing at every record, compute how many records to skip. The customer messages and the 20-position stream are our own example.

• Experience Replay needs a buffer of earlier examples; Reservoir Sampling is one way of deciding what goes into it.
• Continual Learning assumes learning from a stream of data with limited memory, that is, a situation in which not everything can be kept.
• Catastrophic Forgetting [Polski] can get worse when all the examples of an earlier category drop out of a small memory.

The text and illustration were prepared with AI assistance. The message stream and the demonstration are our own simplified educational example; the illustration is a metaphor, not a diagram of the algorithm.