Reservoir Sampling

Um modelo de atendimento ao cliente treina com mensagens que chegam uma após a outra. Você quer guardar algumas delas para repetições futuras, mas só tem quatro vagas na memória. Também não sabe quantas mensagens ainda vão chegar: vinte ou dois milhões.

Ficar com as quatro primeiras? Aí a memória só conhece o começo. Manter sempre as quatro últimas? Aí só conhece o fim. Um sorteio justo de “quatro do total” exigiria guardar tudo ou saber de antemão o número de mensagens. A pergunta é: como, com um número fixo de vagas e um fluxo de comprimento desconhecido, dar a mesma chance a todas as mensagens?

Reservoir Sampling é uma forma de manter uma amostra aleatória de tamanho fixo durante uma única passagem por dados cuja quantidade não conhecemos de antemão. Depois de cada nova mensagem, a memória, aqui chamada de reservatório, contém uma amostra aleatória justa de tudo o que chegou até então. A versão mais simples, o Algorithm R, é descrita por Vitter, Random Sampling with a Reservoir, §2, que a atribui a Alan Waterman.

A receita em dois passos

1. As quatro primeiras mensagens vão para a memória sem sorteio.
2. A mensagem de número $t$ (a quinta, a sexta e todas as seguintes) entra na memória com chance $4/t$. Se entrou, substitui uma das quatro guardadas, escolhida ao acaso.

A quinta mensagem entra, portanto, com chance 4/5; a centésima, com chance 4/100. As mensagens novas têm uma chance cada vez menor de entrar, mas as antigas já correram muitas vezes o risco de ser expulsas, e esses dois efeitos se compensam exatamente. Vamos conferir isso depois da quinta mensagem. Ela entra com chance 4/5. A primeira mensagem só sai quando a quinta entrou (4/5) e caiu justamente na vaga dela (1/4), ou seja, com chance 1/5. Ela fica com chance 4/5: a mesma da quinta. Depois de $t$ mensagens, cada uma delas está na memória com chance $4/t$.

Chance igual não é garantia

No experimento, o fluxo tem 20 mensagens. Dezoito tratam de entrega, e duas, a nona e a décima quinta, tratam de devoluções. Cada execução passa esse fluxo 2.000 vezes e conta com que frequência cada posição ficou na memória no final. Os sorteios são de verdade, então os resultados diferem em uma fração de ponto percentual a cada vez.

Com Reservoir Sampling, cada barra tem cerca de 20%, isto é, 4/20. Ao mesmo tempo, na maioria das execuções nenhuma das duas devoluções fica na memória. O valor exato é 3060/4845, ou seja, cerca de 63%: é essa a parcela dos grupos de quatro, entre todos os que se podem escolher das 20 mensagens, formados apenas pelas dezoito mensagens sobre entrega. O método é justo com cada mensagem, mas não cuida para que toda categoria tenha um representante.

Onde isso faz diferença

No treinamento com repetições, Reservoir Sampling é, às vezes, a forma de gravar no buffer: não é preciso conhecer o comprimento do fluxo nem as fronteiras entre as tarefas. Chaudhry et al., Continual Learning with Tiny Episodic Memories, §2, algoritmo 2 o usam exatamente assim. Nos experimentos deles com tarefas de imagem, ele funcionou bem, exceto com memórias muito pequenas, em que as classes anteriores desapareciam do buffer; nesses casos saíam-se melhor as estratégias que cuidam de uma representação igual das classes (§4.4). É um resultado para classificação de imagens, não para LLMs.

O algoritmo em si não tem a ver com aprendizado de máquina. O trabalho de Vitter, de 1985, descreve o sorteio de registros de uma fita magnética de comprimento desconhecido e propõe variantes mais rápidas, que, em vez de sortear a cada registro, calculam quantos registros pular. As mensagens dos clientes e o fluxo de 20 posições são um exemplo próprio.

• Experience Replay precisa de um buffer com exemplos antigos; Reservoir Sampling é uma das formas de decidir o que entra nele.
• Continual Learning pressupõe o aprendizado a partir de um fluxo de dados com memória limitada, isto é, uma situação em que não dá para guardar tudo.
• Catastrophic Forgetting [Polski] pode se agravar quando todos os exemplos de uma categoria antiga saem de uma memória pequena.

O texto e a ilustração foram preparados com a ajuda de IA. O fluxo de mensagens e a demonstração são um exemplo didático próprio e simplificado; a ilustração é uma metáfora, não um esquema do algoritmo.