RAG Embeddings: Mean Pooling, Memory Bandwidth, and the 12x Batching Speedup | Ameya

Liam, in for Ameya Deshmukh, opens the black box of text embedding. The video shows how per-token vectors are averaged into a chunk vector and why batching 64 chunks beats looping, because memory bandwidth, not math, is the bottleneck.

3:14 video3 min readWatch on YouTube

Turning text into vectors can sound like a black box. In this video, Liam, in for Ameya Deshmukh, argues that it is not, and shows one move, batching, that made the project's embedding step 12 times faster while doing exactly the same math.

One vector per token, not per chunk

Most people get one thing wrong: the transformer does not produce one vector for your chunk. It produces one vector per token. In the video's example, seven tokens go in and seven vectors come out, each 384 numbers wide. The chunk vector you think you have does not exist yet.

Mean pooling

To get a single vector for the chunk, you use mean pooling. The seven token vectors are averaged, element by element, into one. That average is a single 384-dimensional vector, and it is what represents the whole chunk in your index. The video notes that it is simple and is the default for a reason.

Looping versus batching

You could embed one chunk at a time, but almost nobody does. Instead, many chunks are stacked into one matrix and pushed through together. One row becomes 64 rows, so a skinny multiply becomes a fat one, and the same weight matrix serves all of them at once.

The video gives the measurement from the project. Sixty-four separate matrix multiplies took just under 15 milliseconds. One batched multiply of the same 64 took 1.2 milliseconds, a speedup of about 12 and a half times.

Identical arithmetic, different cost

The work was identical. Counting the arithmetic operations both ways gives the same number, 75,497,472, whether looped or batched. So the speedup is not less math, which raises the question of where the time went.

The answer is overhead and memory. Looping pays a fixed setup cost 64 times, while batching pays it once. The weight matrix is loaded from memory once and reused across all 64 rows, instead of being fetched again and again. The bottleneck was never the arithmetic; it was memory bandwidth. The video's analogy is an oven: one cookie or 64, it still takes 12 minutes.

In practice

In practice this is one argument. Sentence Transformers takes a batch size, and handing it your whole list of chunks with a batch size of 64 to 256 lets it do the stacking for you. The video calls this the single biggest win at index time, and the lever that matters most when you index a large corpus.

Scaling to a large corpus

Beyond batching, the video names two more practices:

  • Stream the work in shards so that memory stays bounded.
  • Content-hash each chunk so you only re-embed what actually changed.

Do those three things, it says, and embedding a million chunks stops being scary.

Key takeaways

  • A transformer outputs one vector per token; mean pooling averages them into one chunk vector.
  • In the example, a 384-dimensional vector represents the chunk in the index.
  • Batching 64 chunks took 1.2 ms versus just under 15 ms for 64 separate multiplies.
  • The arithmetic is identical, at 75,497,472 operations either way.
  • The real bottleneck is memory bandwidth and per-call overhead, not math.
  • Use a batch size of 64 to 256, stream in shards, and content-hash chunks.

Try it yourself

Take your own indexer and time encoding your chunks one at a time versus batched at 64. Then confirm that the outputs are identical. Seeing the same vectors come out an order of magnitude faster is the point of the exercise.

Chapters

  1. 0:00Introduction: Demystifying Text Vectorization
  2. 0:40Per-Token Vectors to Chunk Vector: How Mean Pooling Works Single Chunk Encoding vs Matrix Batching
  3. 2:15Performance Benchmark: 15ms Looped vs 1.2ms Batched Matrix Multiply
  4. 3:00Why Batching is 12x Faster: Memory Bandwidth vs Compute Overhead
Full transcript(auto-generated, with timestamps)

Introduction: Demystifying Text Vectorization

[0:00]Namaste circumflex euros, this is Liam in for Amia. Build it with Claude and take it apart. Turning text into vectors sounds like a black box. It isn't. And there's one move a circumflex euros batching a circumflex euros that made the project's embedding step 12 times faster while doing exactly the same math. Let's open the box. Here's the thing most people get wrong. The transformer does not produce one vector for your chunk. It produces one vector per token. Seven tokens go in, seven vectors come out, each 384 numbers wide. The chunk vector you think you have doesn't exist yet. So how do you get a single vector for the chunk? Mean pooling. You average all seven token vectors element by element into one. That average a circumflex euros, a

Per-Token Vectors to Chunk Vector: How Mean Pooling Works Single Chunk Encoding vs Matrix Batching

[0:41]Single 384 dimensional vector a circumflex euros, is what represents the entire chunk in your index. Simple, and it's the default for a reason. Now the speed. You could embed one chunk at a time. Almost nobody does. Instead, you stack many chunks into one matrix and push them through together. One row becomes 64 rows a circumflex euros. A skinny multiply becomes a fat one a circumflex euros. And the same weight matrix serves all of them at once. And here's the measurement from the project. 64 separate matrix multiplies took just under 15 milliseconds. One batched multiply of the same 64 took 1.2. That's a 12 and a half times speed up a circumflex euros for doing the identical work. And I mean identical. Count the arithmetic operations both ways and you get the same number a

Circumflex euros. 75 million 497,472. Exactly the same looped or batched. So the speed up is not less math, which raises the real question. Where did the time go? The answer is overhead and memory. Looping pays a fixed setup cost 64 times. Batching pays it once. And the weight matrix gets loaded from memory once and reused across all 64 rows instead of being fetched again and again. The bottleneck was never the arithmetic a circumflex euros. It was memory bandwidth. It's one oven, one cookie or 64, still 12 minutes. In practice, it's one argument. Sentence Transformers takes a batch size, hand it your whole list of chunks in a batch size of 64 to 256, and it does the stacking for you. This is the single biggest win at index time a circumflex euros, and it's one

Performance Benchmark: 15ms Looped vs 1.2ms Batched Matrix Multiply

[2:15]Keyword. And this is the lever that matters most when you index a large corpus. Batching is the biggest single speed up. Stream the work in shards so memory stays bounded. And content hash each chunk so you only re-embed what actually changed. Do those three and embedding a million chunks stops being scary. So keep three things. The model makes one vector per token. Mean pooling collapses them into one chunk vector. Batching is 12 times faster while doing identical arithmetic because the real bottleneck is memory bandwidth, not math. And at scale batch first a circumflex euros, it's the single biggest index time win, and it costs you one argument. Your turn. Take this prompt to your own indexer. Time encoding your chunks one at a time versus batched at 64 and confirm the outputs are identical.

Why Batching is 12x Faster: Memory Bandwidth vs Compute Overhead

[3:00]When you see the same vectors come out an order of magnitude faster, you'll never loop and encode again. Rag a circumflex euros embeddings and the 12 times win hiding in one argument. Batch it. This is Liam in for Aamia.

More from Humanitarians AI Fellows

Humanitarians AI Lyrical Literacy Project