Recently I played around with some reinforcement learning algorithms and came across the problem of sampling from replay buffers. I thought it is a fun exercise to write a blog post about.

There are two popular strategies to sample from a dataset in reinforcement learning: uniformly, where each data point has the same probability, and prioritised, where each data point has a different probability of getting sampled, i.e. weighted by training error.[^{?cit:replay-buffer}]

The aim of this post is to give an outline how to implement these sampling strategies as efficient algorithms.

# Uniform Distribution

The default strategy in stochastic gradient descent is sampling batches uniformly. All data points have the same probability.

$$
P(x_i) = \frac{1}{n} \; \text{ where $i \in \{0, 1, \dots, n-1\}$}
$$

Here $x_i$ is a data point from our dataset $\mathcal{D} = \{x_0, x_1, \dots, x_{n-1}\}$.

The goal is to sample $k$ items from $\mathcal{D}$. We can either choose to sample without or with replacement. What's the difference?

_Sampling without replacement_ means samples are unique and cannot be sampled twice per batch. On the flip-side, _sampling with replacement_ means items can appear twice in one batch. Implementing the latter is trivial. Simply pick $k$ integers in $\left[0; n\right)$ and you are done.

Sampling without replacement is more interesting. How can we implement this efficiently?

## Sampling without Replacement

Let's think of a physical analogy — picking cards. If you were asked to pick $k$ random cards out of a poker deck without replacement, you might come up with the following strategy: 1. Pick a random card out of one stack of cards, 2. Put that card on a different stack, 3. Repeat.

If we replace stacks with arrays in this description, we already have a quite convincing algorithm, although a few details are missing. Here is one way to do it:

Much like with a stack of cards, we keep two arrays. The first array stores the result (_result array_), the second array stores the integers which we can still choose from (_open array_). Initially, there are zero integers in the result array and $n$ integers in the open array. On each iteration we then randomly pick one integer in the open array, remove it and put it into the result array. After doing this $k$ times, we have a batch of unique samples.

To make this more efficient, we can store both arrays in one "super" array which we divide into two at index $r$. The result array goes from index $0$ to index $r$, while the open array goes from index $r$ to $n-1$ (inclusive). When we want to transfer the sample from open to result array, we simply swap the picked index $i$ with the integer at index $r$ and increase $r$ by one. Figure {?fig:uni-ill} visualises this algorithm.

<Figure caption={

Figure {#fig:uni-ill}: Illustration of uniform algorithm without replacement. First, we pick a random element of the range and swap it with the first element. From now on, ignore the first element and only sample beginning with the second integer. Keep swapping the samples and shrinking the sampling range until we have $k$ integers. The first $k = 3$ integers represent the final sample.

}>

<Include
 src="@/public/uploads/images/sampling-algorithms/sampling-uniform.svg"
 className={`max-w-xl mb-8`}
/>

</Figure>

Let's implement this algorithm using C++. As we need to have an array of size $n$ in memory, it makes sense to reuse that througout multiple samples. Therefore we create a class for the sampler.

<Figure caption={`Figure {#fig:uniform-class}: Class for the uniform sampler. Make sure to free the array in the destructor!`}>

```cpp
#include <stdio.h>
#include <random>

typedef std::mt19937 RandomEngine; // or any other PRNG

class UniformSampler {
public:

 UniformSampler(unsigned int N);
 ~UniformSampler();

 unsigned int *Sample(unsigned int k);

 unsigned int N_;
 unsigned int *array_;

 RandomEngine engine_;

};

UniformSampler::UniformSampler(unsigned int N)
 : N_(N), engine_(RandomEngine { std::random_device{}() }) {
 // Initialize array A to the range [0, N-1)
 array_ = new unsigned int[N_];
 for (int j = 0; j < N_; ++j) {
 array_[j] = j;
 }
}

UniformSampler::~UniformSampler() {
 delete [] array_;
}
```

</Figure>

Now that we have our class, we can implement the algorithm. On each iteration $i$, we sample $j \sim U\{r, N\}$ and swap the $r$-th with $j$-th integer as previously discussed.

<Figure caption={`Figure {#fig:uniform-algo}: Implementation of uniform algorithm.`}>

```cpp
// Sample k integers and return pointer to result array
unsigned int *UniformSampler::Sample(unsigned int k) {
 for (unsigned int r = 0, j; r < k; ++r) {
 // Sample random int in [r; N)
 j = engine_() % (N_ - r) + r;
 // Swap i-th with j-th integer
 std::swap(array_[r], array_[j]);
 }
 // Return pointer to result
 return array_;
}
```

</Figure>

The per-sample runtime of this algorithm is $O(1)$. The memory requirement is $O(n)$. In contrast to other algorithms, the runtime does not depend on $n$. More analysis on the runtime can be found in Figure {?fig:runtime}.

# Unnormalised Categorical Distribution

As mentioned in the introduction, there are other ways to sample from a dataset besides uniform sampling. In the following, I will be focusing on the _categorical distribution_ where each data point has its own probability of getting sampled.

$$
P(x_i) = p_i
$$

Most of the time, however, we do not have probabilities. Instead we are given weights $w_i$ which indicate relative importance of a sample. For example, a set of weights might be $\{1, 3, 8, 1, 3, 2, 1, 4\}$. In this case, the third data point has the highest importance. As a valid probability distribution requires that $\sum_i p_i = 1$, we need to normalise the weights before we can sample.

$$
P(x_i) = \frac{w_i}{\sum_j w_j}
$$

Returning to the example from above, we would have a probability distribution of $\{\frac{1}{23}, \frac{3}{23}, \frac{8}{23}, \frac{1}{23}, \frac{3}{23}, \frac{2}{23}, \frac{1}{23}, \frac{4}{23}\}$. For larger datasets, sampling from this distribution is tricky. Summing over all elements has linear complexity. Moreover, we cannot cache the sum, because in reinforcement learning, for example, weights are changing after each iteration. Is there a way to improve over linear complexity?

## Segment Tree

The answer is yes. We can make use of a binary tree structure to efficiently compute the sum over all elements. The idea is to arrange all data points as leafs of a tree. The key of each leaf is the weight $w_i$ of its data point. To maintain a sum over all weights, the key of a higher-level node is defined as the sum of its child nodes.

$$
\mathrm{key}\left(i\right) = \sum_{j \in \mathrm{children}\left(i\right)} \mathrm{key}\left(j\right)
$$

This is known as a _segment tree_ or _sum heap_. Take a look at Figure {?fig:segment-tree} for an illustration of this data structure.

<Figure caption={

Figure {#fig:segment-tree}: Illustration of a 'sum' [segment tree](https://en.wikipedia.org/wiki/Segment_tree). The leafs of the tree represent the weight of each sample. Every node is equal to the sum of its child nodes. Dividing by the value of the root gives normalised probabilities (shown in green).

}>

<Include
 src="@/public/uploads/images/sampling-algorithms/sampling-categorical.svg"
 className={`mb-8`}
/>

</Figure>

Given a segment tree of the weights, it is straightforward to sample from a categorical distribution. Take a random number $u \in \mathbb{R}$ between $0$ and $1$ sampled from a uniform distribution $u \sim U\{0, 1\}$. Let $v = \hat{w} \cdot u$ where $\hat{w}$ is the value of the root node. Start at the root node and descent left if $v < \mathrm{key}\left(\mathrm{left\,node}\right)$, otherwise set $v = v - \mathrm{key}\left(\mathrm{left\,node}\right)$ and descent right. Once you reach a leaf node, you have your sample.

Essentially, we partition $\left[0, \hat{w}\right)$ into $n$ intervals with the size $w_i$, sample from $\left[0, \hat{w}\right)$ and then find the interval in which that number falls.

To sample without replacement, simply set the weight to zero after each sample and update the tree. This has $\log n$ computational cost. The only drawback is that you need to store the old weights and restore them afterwards, if you want to sample from the same distribution again.

We start implementing this algorithm by designing a class that can hold the tree. To make things easier, we store the tree using an array representation where the two children of node $j$ are at position $2j$ and $2j + 1$. This requires that $n$ is a power of $2$, resulting in an array that has size $2n$.

<Figure caption={

Figure {#fig}: Categorical class. We round up the size of the tree array, so that it is a power of $2$.

}>

```cpp
#include <cassert>
#include <stdio.h>
#include <random>

typedef std::mt19937 RandomEngine;

class CategoricalSampler {
public:

 CategoricalSampler(unsigned int N);
 ~CategoricalSampler();

 void SampleWithReplacement(unsigned int k, unsigned int *result);
 void SampleWithoutReplacement(unsigned int k, unsigned int *result);

 void SetWeight(unsigned int j, float weight);
 const float& GetWeight(unsigned int j);

 unsigned int N_; // Dataset size
 unsigned int D_; // Depth of tree
 float *w_; // Tree array
 RandomEngine engine_;

};

// From https://stackoverflow.com/a/11376759
#define LOG2(X) ((unsigned) \
 (8*sizeof (unsigned long long) - __builtin_clzll((X)) - 1))

CategoricalSampler::CategoricalSampler(unsigned int N)
 : N_(N), engine_(RandomEngine { std::random_device{}() }) {
 // Calculate tree depth, round up
 D_ = LOG2(N - 1) + 1;
 assert(D_ < 32); // Would cause overflow
 // Create tree structure using array representation
 unsigned int N1 = 1 << D_; // Index of first leaf node, 2^(D-1)
 unsigned int N2 = 2 << D_; // Length of array, 2^D
 w_ = new float[N2];
 // Initialize the weights to w_i = 1
 std::fill(w_ + N1, w_ + N1 + N, 1.0f);
 // Set weight of all excess nodes to zero
 std::fill(w_ + N1 + N, w_ + N2, 0.0f);
 // Heapify
 for (unsigned int i = N1 - 1; i > 0; i -= 1) {
 w_[i] = w_[2 * i] + w_[2 * i + 1];
 }
}

CategoricalSampler::~CategoricalSampler() {
 delete [] w_;
}
```

</Figure>

Next, we implement some basic functions to maintain the sum tree. `Heapify` walks up the tree and restores the sum condition. `SetWeight` and `GetWeight` are for users who wish to update the distribution.

<Figure caption={`Figure {#fig}`}>

```cpp
inline void Heapify(float *w, int j) {
 // Walk up tree and restore sum condition
 while (j > 1) {
 j /= 2;
 w[j] = w[2 * j] + w[2 * j + 1];
 }
}

void CategoricalSampler::SetWeight(unsigned int i, float weight) {
 unsigned int j = (1 << D_) + i; // Position of w_j in array
 w_[j] = weight;
 Heapify(w_, j); // Maintain sum condition
}

const float& CategoricalSampler::GetWeight(unsigned int i); {
 return w_[(1 << D_) + i];
}
```

</Figure>

After having all the setup done, we can focus on implementing the sampling algorithm as described earlier.

<Figure caption={

Figure {#fig}: Sampling algorithm for a categorical distribution, without replacement.

}>

```cpp
void CategoricalSampler::
 SampleWithoutReplacement(unsigned int k, unsigned int *result) {
 const unsigned int N1 = 1 << D_;
 float w_old[k]; // Store old weights here
 float v; // Sampled value
 for (unsigned int i = 0, j; i < k; i += 1) {
 // Pick random number between [0; w_[1])
 v = std::uniform_real_distribution<>(0.0f, w_[1])(engine_);
 // Binary search through segment tree
 j = 1;
 do {
 j *= 2;
 if (w_[j] < v) {
 v -= w_[j];
 j += 1;
 }
 } while (j < N1);
 // Set output index
 result[i] = j - N1;
 // Store weight
 w_old[i] = w_[j];
 // Remove weight from tree
 v = w_[j];
 while (j > 0) {
 w_[j] -= v;
 j /= 2;
 }
 }
 // Restore original state
 for (unsigned int i = 0, j; i < k; i += 1) {
 j = N1 + result[i];
 v = w_old[i];
 while (j > 0) {
 w_[j] += v;
 j /= 2;
 }
 }
}
```

</Figure>

<More className="my-8" teaser="code for sampling with replacement">

<Figure className="lf" caption={

Figure {#fig}: We can also implement a function that samples with replacement.

}>

```cpp
void CategoricalSampler::
 SampleWithReplacement(unsigned int k, unsigned int *result) {
 unsigned int N1 = 1 << D_; // 2^{d-1}
 float v; // sampled value

 for (unsigned int i = 0, j, l; i < k; i += 1) {
 // Pick random number between [0; w_[1])
 v = std::uniform_real_distribution<>(0.0f, w_[1])(engine_);
 // Binary search through segment tree
 j = 1;
 l = D_;
 do {
 j *= 2;
 l -= 1;
 if (w_[j] < v) {
 v -= w_[j];
 j += 1;
 }
 } while (l);
 // Set output index
 result[i] = j - N1;
 }
}
```

</Figure>

</More>

This segment tree-based algorithm has runtime complexity $O(\log n)$ per sample. Its memory complexity is $O(n + k)$.

# Runtime Analysis

You might wonder how well these algorithms perform for bigger datasets in practice. To find that out, I ran some performance tests. Take a look at Figure {?fig:runtime}.

<Expose className="my-10">

<Figure className="col-text">
 <Image
 src="/uploads/images/sampling-algorithms/runtime_analysis.png"
 alt="Runtime analysis"
 width={700}
 height={500}
 contain={true}
 responsive
 />
</Figure>

<figcaption className="col-text pt-2 mt-4 text-gray-400">

Figure {#fig:runtime}: Runtime analysis of the presented algorithms. (1) Runtime of categorical sampling shows log-complexity with respect to $N$, and $k$ fixed to $1024$. (2) Runtime of categorical sampling has linear complexity with respect to $k$, and $N$ fixed to $64 000$. (3) Runtime of uniform algorithm with respect to $k$, and $N$ fixed to $64 000$. (4) Complexity of categorical over uniform distribution, with $N = 64000$.

</figcaption>

</Expose>

The log-complexity of categorical sampling w.r.t. to $n$ is noticeable in plot (1). There is some performance trade-off when using sampling without replacement. It takes around $1.5$ms to sample a batch of size $1024$ from a dataset with $100$ million data points (which by today's machine learning standards is a reasonably large dataset). Categorical sampling is around $25$ to $30$ times slower compared to uniform sampling for $N = 64000$, as seen in plot (4).

# Effect on Training Neural Nets

We discussed sampling _without replacement_ throughout this post. It might not be obvious what the effect on training a neural net is. From a theoretical standpoint, it reduces the variance of an estimator. Take for example a mean estimator $Y = \frac{1}{k}\sum_{i=1}^k X_i$ with $X_i \sim \mathcal{D}$. Using sampling without replacement, the variance of this estimator reduces to[^{?cit:utexas}]

$$
\hat{\sigma}^2 = \frac{N-k}{N-1} \cdot \sigma^2
$$

where $\sigma^2$ is the variance of the same estimator but _with replacement_ sampling.

In domains where $N$ is small, this can make a difference. If $N$ is large, however, the factor $\frac{N-k}{N-1}$ gets closer and closer to $1$. For example, with a batch size of $b = 1\%$ of the dataset size, we get a factor of $N \cdot \frac{1 - b}{N-1} \approx 1 - b = 99 \%$. Sampling without replacement reduces the variance by $1\%$ in this case.

Although this applies only to mean estimators, the analogy might give you some intuition about the effect on neural network training of sampling without replacement.

# Conclusion

In this post, I presented algorithms to sample from a dataset uniformly or weighted. These algorithms find application in reinforcement learning for prioritised replay buffers [^{?cit:replay-buffer}] or in supervised learning when sampling batches. The effect of sampling without replacement is larger for small datasets.

Further research into the effect of prioritising training examples based on learning error would be interesting. The algorithms presented here could be used as a basis for ideas in that space.

---

If you want to read more about sampling algorithms, take a look at Tim Vieira's blog: https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/

[^{#cit:replay-buffer}]: https://arxiv.org/abs/1511.05952 with a good breakdown found at https://danieltakeshi.github.io/2019/07/14/per/
[^{#cit:utexas}]: https://web.ma.utexas.edu/users/parker/sampling/woreplshort.htm
[^more]:

Sampling Algorithms

South Africa is, I would say, one of the most beautiful and at the same time heart-breaking
countries. Its beautiful nature and devastating townships are in stark contrast.

This year I made it back to the country that has all of the big five — lion, leopard, elephant,
rhino, and buffalo. It was a stunning experience.

We started our trip in Johannesburg, going to the north-west to the Madikwe Game Reserve.

<Figure caption={`Overlooking Madikwe Game Reserve`}>
 <Image
 src="/uploads/images/safari-in-south-africa/madikwe.jpg"
 alt="Overview of Madikwe Game Reserve"
 width={600}
 height={400}
 priority
 />
</Figure>



<Expose>

<Figure className="col-extend" full>
 <Image
 src="/uploads/images/safari-in-south-africa/elephant-2.jpg"
 alt="Elephant on safari"
 width={1200}
 height={900}
 priority
 />
</Figure>

<Figure className="col-start-4 col-span-6">
 <Image
 src="/uploads/images/safari-in-south-africa/early-morning-coffee.jpg"
 alt="Drinking coffee early in the morning in the safari park"
 width={300}
 height={400}
 />
</Figure>

<Figure className="col-start-11 col-span-8 justify-self-stretch place-self-end">
 <Image
 src="/uploads/images/safari-in-south-africa/tree.jpg"
 alt="Magical tree"
 width={400}
 height={300}
 />
</Figure>

</Expose>



<Figure>
 <Image
 src="/uploads/images/safari-in-south-africa/lion-1.jpg"
 alt="A lion"
 width={600}
 height={450}
 />
</Figure>

<Figure>
 <Image
 src="/uploads/images/safari-in-south-africa/safari-3.jpg"
 alt="Sun rising over Madikwe"
 width={1200}
 height={900}
 />
</Figure>

<FloatFigure expose="col-extend" full>
 <Image
 src="/uploads/images/safari-in-south-africa/lion-3.jpg"
 alt="Lion that is photographed by people out of safari vehicles"
 width={600}
 height={450}
 />
</FloatFigure>

<Figure>
 <Image
 src="/uploads/images/safari-in-south-africa/safari-4.jpg"
 alt="Picture out of our car into the sun"
 width={600}
 height={800}
 />
</Figure>

<FloatFigure expose="col-extend" full caption={`Pink and White Skies`}>
 <Image
 src="/uploads/images/safari-in-south-africa/dusk.jpg"
 alt="Artsy picture of dusk in the desert"
 width={600}
 height={450}
 />
</FloatFigure>

Waking up at 5 in the morning while the sun was still down, coffee at 8. Our schedule was insane.
The early bird, gets the worm – I guess; getting up early meant we were getting all the
sunrises. At 9 we were done. We went back to sleep, only to be ready for the next safari at 4 in
the afternoon.

After three awesome days in Madikwe, it was about time to move to our next destination. We went
to the Welgevonden Reserve. It was north of Madikwe, near the border to Botswana. Getting there
was a bumpy ride. Roads on a map do not necessarily translate to tarred streets in South Africa.
Take note.



<Figure
 className="mx-auto"
 intrinsic
 caption={`The three stops on our journey: Madikwe, Welgevonden, Cape Town`}
>
 <Image
 src="/uploads/images/safari-in-south-africa/maps.png"
 alt="A map of South Africa which highlights our journey"
 className="rounded-xl"
 width={300}
 height={300}
 intrinsic
 />
</Figure>



Once we arrived, everything changed. It was quite interesting how much the landscape differed.
Even in Madikwe, on the same day, it felt like we were visiting different locations on earth
although it was all part of the same game reserve. Welgevonden felt completely different yet
again. If Madikwe was like a steppe, this game reserve was like a jungle.

<FloatFigure expose="col-extend" full caption={`Do you feel the jungle?`}>
 <Image
 src="/uploads/images/safari-in-south-africa/jungle.jpg"
 alt="Sun rising in the Welgevonden Reserve"
 width={1200}
 height={900}
 />
</FloatFigure>

<Figure>
 <Image
 src="/uploads/images/safari-in-south-africa/safari-1.jpg"
 alt="Elephant on the street"
 width={600}
 height={800}
 />
</Figure>

<Figure>
 <Image
 src="/uploads/images/safari-in-south-africa/giraffe.jpg"
 alt="Giraffe in the wild"
 width={600}
 height={450}
 />
</Figure>

<Figure>
 <Image
 src="/uploads/images/safari-in-south-africa/elephant-1.jpg"
 alt="Elephant moving away from our car"
 width={600}
 height={450}
 />
</Figure>

In this game reserve the elephants were messing with us. Usually, they came running towards our
safari vehicle when we encountered them. Get's your blood pumping. Apparently they were in their
musth phase which means they are more aggressive. With their 6 tons of weight, elephants could
easily turn around a safari vehicle.

After a few days and many elephant encounters, we drove back to Johannesburg. From there we flew
to our last destination: Cape Town, one of my favorite cities I have been to. We stayed in Cape
Town over New Years Eve. It’s crazy how cheap good food is in this city. One restaurant I can
highly recommend is called Manna Epicure in Kloof Street. It's damn good. And also, did I tell
you about the light?



<Expose>

<Figure
 caption={`Mystic haze`}
 className="col-start-4 col-span-24 md:col-start-4 md:col-span-12"
>
 <Image
 src="/uploads/images/safari-in-south-africa/from-table-mountain.jpg"
 alt="Clouds over Cape Town and Lions Head"
 width={800}
 height={600}
 />
</Figure>

<Figure
 caption={`Lions Head at dusk`}
 className="col-start-12 col-span-16 md:col-start-18 md:col-span-8 md:row-span-2 
 justify-self-stretch place-self-center"
>
 <Image
 src="/uploads/images/safari-in-south-africa/lions-head.jpg"
 alt="Beautiful lions head during dusk"
 width={600}
 height={800}
 />
</Figure>

<Figure
 caption={`Clouds over the twelve apostles`}
 className="col-start-4 col-span-16 md:col-start-4 md:col-span-8 "
>
 <Image
 src="/uploads/images/safari-in-south-africa/twelve-apostels.jpg"
 alt="Clouds over twelve apostels"
 width={600}
 height={800}
 />
</Figure>

</Expose>

Safari in South Africa


The best startups are founded by ingenius people. I agree with this. But there is another, somewhat more overlooked aspect to great founders. The best founders live in the future.

In order to find problems that are yet to be solved, the future is the best place to be. Only when you take for granted what is novel to other people, you are able to assemble the puzzle and look for missing pieces.
Mark Zuckerberg did not create Facebook because he looked for good ideas to pursue. He already spent so much of his time on the internet, more than most people at that time. For him it was only natural to share things with his friends online. That’s why he build Facebook. He lived in the future.


Live in the Future


Imagine a world where artificial general intelligence is a thing, where robots assemble robots, where computers cure cancer, and where everything we do today is automated. In such a robot world, the problems arise with humans – not robots. What is left to do for us, if we are replaced in everything we do today? More importantly, how can computers align with our values, so they work with us and not against us?

## The Solution

In my opinion, an inevitable solution is a device called neural lace. It’s a brain implant that extends your intelligence by connecting neurons in your brain to an artificial neural network. It works by replicating processes of your brain, with the speed and scalability perks of silicon. Wearing such a device must feel like an epiphany, where suddenly you can make sense of things you could not before.

Why is a neural lace inevitable? First of all, it keeps humans at work – as they are not replaced but augmented. Second, it’s of high economic interest because it increases productivity. Finally, it solves the problem of value-alignment: a neural lace doesn’t act on its own, keeping humans in charge. These reasons make me to strongly believe a neural lace to be inevitable in a post-AGI world.
Although a neural lace may solve some of AGI’s problems, it inherently comes with new, bigger controversies.

## Bye, Meritocracy

Top of society has always consisted of the brightest minds, the greatest thinkers, the ones who come up with best ideas and most ingenious inventions. If we assume a neural lace to be inevitable, people having access to such a device will outsmart others. They will have unattainable brain capabilities, will be able to memorize and comprehend more than their peers and will be able to see connections where others don’t. Hence, access to such devices will be very desirable. To be at top of the society will mean having a neural lace. This slashes even the slightest delusion of a meritocracy, as only the one percent may be able to afford such devices. Suddenly being rich means being smart.

As cognitive abilities grow, leadership and value-creation in the society will shift towards intelligence-augmented people. Being 100x more intelligent than your rivals will obviously mean outperforming them. Being smarter also means producing things, that people need, more efficiently. As more value is created by intelligence augmented people this way, a social divide is unavoidable. If 99% of the population accounts only for 1% of value creation in the economy, society will be split. Two groups emerge: a small, elite group of people who have access to neural laces, and the rest. In my opinion, this will be the biggest controversy in the history of mankind.


Recent posts