outlines icon indicating copy to clipboard operation
outlines copied to clipboard

Add Particle Filter

Open rlouf opened this issue 1 year ago • 1 comments

See Nicolas Chopin's book for a really nice introduction to the topic. The algorithms consists in carrying $N$ particles for each sequence in the batch, and at each step to:

  1. Sample a new token for each particle using the next-token logits.
  2. Resample the particles.

We use the multinomial resampling function in this first PR, although it is known to have very large variance. To make the implementation easier we combine (1) and (2) in a single step, similarly to what we do with beam search.

Note that there is a subtlety when doing structured generation. We can think of the simple following scheme to sample from the distribution of sequences that follow the structure:

  1. Move particles by one step using the unbiased next-token logits;
  2. Set the weight of each invalid particle to $-\infty$
  3. Resample

But this can be very inefficient. Instead, we move particles using a specific proposal: using the biased next-token logits. Since this is not exactly sampling from the original distribution we need to resample the particles using the factor $P_i / \tilde{P}_i$ as a weight where $P_i$ is the unbiased probability of token $i$ and $P_i$ the biased probability of token $i$ (importance sampling).

Note: I am wondering if we should correct the Beam Search algorithm as well.

rlouf avatar Feb 16 '24 10:02 rlouf

Doing this I started to wonder if we shouldn't see and implement greedy and multinomial sampling as particular cases of more general samplers (resp. a form of beam search and a form of particle filtering).

rlouf avatar Feb 29 '24 12:02 rlouf