Mistral 7B & Mixture of experts(8x7B)

Lets now try to see MOE work we start with explaining the Mistral 7B paper and then see MOE and different variants of it from dense to sparce, different sampling techniques.

Its a decoder only model which its architecture which comes with two flavors, one with normal feed forward layer and using smaller expert FFN layers which replace the single layer processing all tokens.

Sliding window attention

In previous works where we apply only masked attention where we mask future tokens in order to give our model causality, we’re seeing the whole k-1 tokens while computing the self attention. The idea of windowed attention stems from convolutional works and it states that we should only pay attention to what’s in specific window. So if we say our window size is W, we only attend to the previous w-1 tokens.

While it seems like we’re losing information by not paying attention to the rest of the tokens with the name of efficiency, since we’ve multiple layers, the attention will still be there but becoming indirect connection.

As we concatenate different layers, on the previous layer, each part of the token t was dependent on the previous (w-1) tokens where w is the size of the window.

Rolling Buffer cache

As we’ve discussed about KV cache where instead of calculating the attention(Q*K^T) * V) every time while the only thing that’s actually changing is the query. We store the KV of previous runs and only append new ones. This work instead of storing every KV cache, we only store the cache of the previous W tokens because we’ll not be computing attention to the whole token space.

The way to do this rolling process is until the cache is full append the tokens by moving the write pointer to right. Once the cache is full we point the write pointer to the starting position and replace it with the latest token. But, this will create order mismatch so we fix the starting pointer and move the rest of the tokens to the left of it so the write pointer appears at latest. This process is called unrolling.

Chunking and prefilling

Since the size of our kv cache is limited and also we get the user input, we can divide our user input into chunks and add them into our kv cache after computing the attention.

Lets say our chunk currently contains the first four words and our KV-cache is empty. We compute the attention with these chunk of our attention and add it to our cache.

When we process the next chunk, we take our query and multiply it with our current query + previous KV cache. This will allow us to pay attention to our previous close related input while keeping our attention span limited.

Model Sharding

This is a technique we use when our model is to big to fit into a single GPU and instead we do divide our layers into groups and feed each groups into our GPU. While inferencing we take the output of the previous layers(in previous GPU) and add it into the next layer in the next GPU.

But this means we will be utilizing one gpu at a time making the others idle or creating bubles.

To avoid bubbles we can divide our model into batches and process a batch at a time and when one batch gets transferred to the next gpu it will take on the next batch giving it much distributed training and use more resources. If we still have bubbles we can increase our batch size and process more batches.

Mixture of Experts

It’s a process of replacing our FFN layer with small expert networks which are trained on subset of the db. router is a ffn layer which given experts and token embeddings it determine the probability of each experts processing this token and it will be used to give weight to the output of these experts. This score is called token-expert affinity score and will be used in weighted sum.

Dense MOE

This is the simplest case of MOE, where we process the token with each experts and multiply the embedding with the probability distribution we got from our router network. While this is simple, it comes at the cost of same computational usage with one giant network as we’re still using the full network. Dense MOE are generally more expressive than the vanilla FFN as they can have more parameters and they get processed by a specific network specializing on them.

Sparce MOE

Since the dense one doesn’t actually reduce any computational burden, we can take the top k experts and process the tokens making our network only compute part of the network while the low probability experts stay untouched. The simplest form of this is top-1, where we take the arg max of the expert affinity score which will give us back the network having the top probability

But this top one approach causes router collapse, where it choses one expert over and over again for all sort of tokens leaving the other network idle.

So we can increase the k to process our tokens by multiple experts and take the weighted sum of it. There are several ways to do this, by taking the top 2, by taking the argmax for one and selecting the other one randomly.

References

• mistral.ai/news/mixtral-of-experts
• Umar Jamil - Mistral / Mixtral Explained: Sliding Window Attention, Sparse Mixture of Experts, Rolling Buffer
• Neural Breakdown with AVB - Writing Mixture of Experts LLMs from Scratch in PyTorch