Example 1
\[
\mathbb{E}_f[w(x)]
=
\frac{1}{n}\sum_{i=1}^{n}w(x_i),
\qquad x_i\sim f(x)
\]
but we cannot sample from \(f(x)\).
Importance Sampling
Aditya Desai
but we cannot sample from \(f(x)\).
Sometimes we may want to purposefully choose a different distribution,
even if we can sample from \(f(x)\).
attention output \(=a_1v_1+a_2v_2+\cdots+a_nv_n,\quad v_i\in\mathbb{R}^d\)
goal is to sample and estimate \(S\)
where
Q. What are the cases in which this is not desirable sampling
Example: values of \(x_i\)
but variance will be high?
Q. What should be a good distribution?
→ A distribution that samples \(x_1\) with high probability (very close to 1)
* This condition is actually the case in attention computation.
depending on dot products subsequent keys far from query \(q\) get exponentially small!
→ Random Sampling will not work!
Formalise this idea in importance Sampling.
Estimator (left-hand side labeled \(\hat\mu\)):
* Unbiased?
* Unbiased? Yes.
So this works for any proposal \(g(x)\)
What's a good \(g(x)\)?
→ One with least variance!!
Problem: Which \(g\) minimises this?
Choose a \(g(x)\) which minimises this.
Simplicity Consider discrete case
rewrite:
If every "\(i\)" is a dimension
Cauchy Schwartz
There is a integral version of this
minimum value
is Cauchy Schwartz tight?
Yes!
When is minimum achieved?
This is the best proposal distribution.
Does it make sense?
\(|w(x)\cdot f(x)|\) is exactly the "contribution"
of \(i\)th term to the expectation.
\(w(x)>0\ \forall x\), then what is the variance?
Something Wrong?
No. \(g(x)\) needs you to know \(\mathbb{E}_f(w(x))\)!
In practice, we will be trying to cheaply get a proposal \(g(x)\)
that is close to \(f(x)|w(x)|\).
How should we sample terms?
Best distribution \(=\)
\(a_i\) computation will need entire KV Cache reading.
→ What is a good \(g(x)\)?
[Think] that is cheaply computable
we will answer this question partially in the Seminar
and hopefully in some future class.
Rejection Sampling
Aditya Desai
So importance sampling lets us compute \(\mathbb{E}_f[w(x)]\)
but what if we want to get samples from \(f(x)\)
but cannot sample directly from it.
Example: Get a point \(x\) in 2D
sampled uniformly at random from a circle of radius \(r\)
using only \(U[0,1]\) variables.
* Sample uniformly in square
→ Reject if outside the circle
Now consider the generic case
\(f(x)\): target distribution
\(g(x)\): distribution to be used.
How do we make sure that probability of getting \(x\) is actually \(f(x)\)?
if \(\underline{g(a) > f(a)}\)
→ Reduce the probability.
→ start rejecting the sample!
Sample \(a \sim g(x)\)
if \(g(a) \ge f(a)\):
\(r = \mathrm{Uniform}[0,1]\)
if \(r > f(a)/g(a)\)
Reject and go sample again
otherwise
Accept, return \(a\)
else ??
What about if \(f(a) > g(a)\)
→ we need to increase the probability of \(a\)
But how?
⇒ Decrease probability of other "\(x\)" enough so that we can relatively use \(g(a)\)
Envelope: Choose an \(M\) s.t. \(Mg(x) > f(x)\ \forall x\)
Algo:
Sample \(a \sim g(x)\)
\(r = U[0,1]\)
if \(r < \dfrac{f(a)}{Mg(a)}\): Accept
else: Reject and Sample again
Example of where to use the Rejection Sampling (Cooked up)
Consider Attention Computation.
Let us say you wanted to approximate the denominator with a sample
What is the best sampling probability distribution?
→ Importance Sampling say
Note \(q\) is not available beforehand in decode.
Say somehow we have obtained the sample proportional to
where \(\bar q\) is some "representative" query.
Now can you get a sample from this subset that is s.t.
where \(q\) is the actual query?
→ Can use Rejection Sampling!!
Reservoir Sampling
Aditya Desai
Previous Cooked up example:
Assumed that we had a sample of KV
But in a auto-regressive generation total KV tokens keep on increasing
Q. How do we "maintain" a sample in that case?
A. Reservoir Sampling.
Consider a stream of data
at each point \(t\), we want to keep a sample of size \(k\)
such that
"Stream": we process the data point only once and "throw" it away.
At time \(t\) we look at \(x_t\)
"Build a streaming algorithm for sampling"
let us say we have a sample
At time \((t+1)\), we got \(x_{t+1}\)
Q. Choose \(x_{t+1}\) with Prob?
Q. If chosen, throw one of the other sampled data with uniform probability
Thus the algorithm works.
Another simple algorithm: Random Key algorithm.
→ Every \(x_i\) is assigned a random value \(r_i\sim U[0,1]\)
→ At each point \(t\), the sample is the \(x_i\)s with
Highest \(k\) values.
→ random key assignment gives a permutation.
→ probability that any key is chosen is \(k/t\) at all times \(t\).
→ At each time you only have to keep track of top \(k\) \(x_i\)s and their scores
If \(r_{t+1} \lt \min R\):
discard \(x_{t+1}\)
Else
Add \((x_{t+1},r_{t+1})\) and discard the min
Weighted Reservoir Sampling
Aditya Desai
Stream
At time \(t\gt k\), \(k\) sample size.
WRS without "Replacement" \([P\text{-}WRS\text{-}w/o]\)
Maintain a sample \(k\). s.t at each point in time \(t\gt k\), the probabilities of inclusion match the following procedure
"If you construct \(k\) sample iteratively the probability that \(i\) is sampled at step \(t\) is \(\propto w_i\)"
Note that under this the marginal probabilities of \(x_i\) in \(k\)-sample is NOT proportional to \(w_i\)
eg. \(\{w_1,w_2,w_3\}\), \(k=2\)
How do we design a random key algorithm?
Sort \(X_i\) and pick top-\(k\).
so that
How should we pick \(\mathcal{F}\)??
Simplicity: let us only look at \(k=1\)
I want probability that \(i\) lies at right most to be \(\propto w_i\)
What should density of \(X\) look like??
If \(w_1\gt w_2\gt w_3\) then \(\leftarrow\)
do we know such \(g(x)\)?
Then density should be such that
Thus "Potentially" \(\mathcal{F}(r_i)=r_i^{1/w_i}\) can work.
→ We have not proved anything yet. We have just come up with "a potential" proposal that "might" work.
Proof: Full proof is out of scope. But we can prove a lemma.
lemma:
If \(r_1\) and \(r_2\) are \(U[0,1]\),
Proof:
Can you extend "non-random-key" algo to WRS. [Exercise]