SeeMoE: Implementing a MoE Vision Language Model from Scratch

<p>

TL;DR: In  this blog I implement a mixture of experts vision language model consisting of an image encoder, a multimodal projection module and a mixture of experts decoder language model in pure pytorch. Thus, the resulting implementation could be thought of as a scaled down version of Grok 1.5 Vision and GPT-4 Vision (both have vision encoders connected to a MoE Decoder model via a projection module). The name ‘seeMoE’ is my way of paying homage to Andrej Karpathy’s project ‘makemore’ because for the decoder used here I implement a character level autoregressive language model much like in his nanoGPT/ makemore implementation but with a twist. The twist being that it's a mixture of experts Decoder (much like DBRX, Mixtral and Grok). My goal is for you to have an intuitive understanding of how this seemingly state of the art implementation works so that you can improve upon it or use the key takeaways to build more useful systems.</p>
<p>The entire implementation can be found in seeMoE_from_Scratch.ipynb in the following repo: <a href="https://github.com/AviSoori1x/seemore" rel="nofollow">https://github.com/AviSoori1x/seemore</a></p>
<div align="center">
<img alt="seemore" height="500" src="https://github.com/AviSoori1x/seemore/blob/main/images/seeMoE.png?raw=true" width="500"/>
</div>
<p>If you've read my other blogs on implementing mixture of experts LLMs from scratch: <a href="https://huggingface.co/blog/AviSoori1x/makemoe-from-scratch">https://huggingface.co/blog/AviSoori1x/makemoe-from-scratch</a> and implementing a vision language model from scratch: <a href="https://huggingface.co/blog/AviSoori1x/seemore-vision-language-model">https://huggingface.co/blog/AviSoori1x/seemore-vision-language-model</a>, you'll realize that I'm combining the two to implement seeMoE. Essentially, all I'm doing here is replacing the feed-forward neural network in each transformer block of the decoder with a mixture of experts module with noisy Top-K gating. More information on how this is implemented, is given here: <a href="https://huggingface.co/blog/AviSoori1x/makemoe-from-scratch">https://huggingface.co/blog/AviSoori1x/makemoe-from-scratch</a>. I strongly encourage you to read these two blogs and carefully go through the repos linked to both the blogs before diving into this.</p>
<p>In ‘seeMoE’, my simple implementation of a Mixture of Experts vision language model (VLM), there are 3 main components.</p>
<div align="center">
<img alt="seemore" height="500" src="https://github.com/AviSoori1x/seemore/blob/main/images/moevlm.png?raw=true" width="500"/>
</div>
<ul>
<li><p>Image Encoder to extract visual features from images. In this case I use a from scratch implementation of the original vision transformer used in CLIP. This is actually a popular choice in many modern VLMs. The one notable exception is Fuyu series of models from Adept, that passes the patchified images directly to the projection layer.</p>
</li>
<li><p>Vision-Language Projector - Image embeddings are not of the same shape as text embeddings used by the decoder. So we need to ‘project’ i.e. change dimensionality of image features extracted by the image encoder to match what’s observed in the text embedding space. So image features become ‘visual tokens’ for the decoder. This could be a single layer or an MLP. I’ve used an MLP because it’s worth showing.</p>
</li>
<li><p>A decoder only language model with the mixture of experts architecture. This is the component that ultimately generates text. In my implementation I’ve deviated from what you see in LLaVA a bit by incorporating the projection module to my decoder. Typically this is not observed, and you leave the architecture of the decoder (which is usually an already pretrained model) untouched. The biggest change here is that, as mentioned before, the feed forward neural net / MLP in each transformer block is replaced by a mixture of experts block with a noisy top-k gating mechanism. Basically each token (text tokens + visual tokens that have been mapped to the same embedding space as text tokens) are only processed by top-k of the n experts in each transformer block. So if it's a MoE architecture with 8 experts and top 2 gating, only 2 of the experts will be activated.</p>
</li>
</ul>
<p>Since the Image Encoder and Vision Language Projector are unchanged from seemore (Linked above. Repo here: <a href="https://github.com/AviSoori1x/seemore" rel="nofollow">https://github.com/AviSoori1x/seemore</a>), I encourage you to read the blog/ go through the notebooks for details on those. </p>
<p>Now let's revisit the components of a sparse mixture of experts module: </p>
<ol>
<li>Experts - just n vanilla MLPs</li>
<li>A gating/ routing mechanism</li>
<li>weighted summation of activated experts based on the routing mechanism</li>
</ol>
<img alt="seemore" height="700" src="https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/experts.png" width="700"/>
<p>First, the 'Expert' which is just an MLP like we saw earlier when implementing the Encoder. </p>
<pre><code class="language-python"><span class="hljs-comment">#Expert module</span>
<span class="hljs-keyword">class</span> <span class="hljs-title class_">Expert</span>(nn.Module):
    <span class="hljs-string">""" An MLP is a simple linear layer followed by a non-linearity i.e. each Expert """</span>

    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, n_embed</span>):
        <span class="hljs-built_in">super</span>().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embed, <span class="hljs-number">4</span> * n_embed),
            nn.ReLU(),
            nn.Linear(<span class="hljs-number">4</span> * n_embed, n_embed),
            nn.Dropout(dropout),
        )

    <span class="hljs-keyword">def</span> <span class="hljs-title function_">forward</span>(<span class="hljs-params">self, x</span>):
        <span class="hljs-keyword">return</span> self.net(x)
</code></pre>
<p>The routing module decides which experts will be activated. Noisy top k gating/ routing adds a bit of gaussian noise to ensure that there's a fine balance between exploration and exploitation in picking the top-k experts for each token. This reduces the odds of the same n experts getting picked everytime, which defeats the purpose of having a larger parameter count with sparse activation for better generalizability.</p>
<img alt="seemore" height="700" src="https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/noisytopkgating.png" width="700"/>
<pre><code class="language-python">
<span class="hljs-comment">#noisy top-k gating</span>
<span class="hljs-keyword">class</span> <span class="hljs-title class_">NoisyTopkRouter</span>(nn.Module):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, n_embed, num_experts, top_k</span>):
        <span class="hljs-built_in">super</span>(NoisyTopkRouter, self).__init__()
        self.top_k = top_k
        <span class="hljs-comment">#layer for router logits</span>
        self.topkroute_linear = nn.Linear(n_embed, num_experts)
        self.noise_linear =nn.Linear(n_embed, num_experts)

    
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">forward</span>(<span class="hljs-params">self, mh_output</span>):
        <span class="hljs-comment"># mh_ouput is the output tensor from multihead self attention block</span>
        logits = self.topkroute_linear(mh_output)

        <span class="hljs-comment">#Noise logits</span>
        noise_logits = self.noise_linear(mh_output)

        <span class="hljs-comment">#Adding scaled unit gaussian noise to the logits</span>
        noise = torch.randn_like(logits)*F.softplus(noise_logits)
        noisy_logits = logits + noise

        top_k_logits, indices = noisy_logits.topk(self.top_k, dim=-<span class="hljs-number">1</span>)
        zeros = torch.full_like(noisy_logits, <span class="hljs-built_in">float</span>(<span class="hljs-string">'-inf'</span>))
        sparse_logits = zeros.scatter(-<span class="hljs-number">1</span>, indices, top_k_logits)
        router_output = F.softmax(sparse_logits, dim=-<span class="hljs-number">1</span>)
        <span class="hljs-keyword">return</span> router_output, indices
</code></pre>
<p>Now both noisy-top-k gating and experts can be combined to create a sparse Mixture of Experts module. Note that weighted summation calculation has been incorporated to yield the output for each token in the forward pass.</p>
<pre><code class="language-python"><span class="hljs-comment">#Now create the sparse mixture of experts module</span>
<span class="hljs-keyword">class</span> <span class="hljs-title class_">SparseMoE</span>(nn.Module):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, n_embed, num_experts, top_k</span>):
        <span class="hljs-built_in">super</span>(SparseMoE, self).__init__()
        self.router = NoisyTopkRouter(n_embed, num_experts, top_k)
        self.experts = nn.ModuleList([Expert(n_embed) <span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(num_experts)])
        self.top_k = top_k

    <span class="hljs-keyword">def</span> <span class="hljs-title function_">forward</span>(<span class="hljs-params">self, x</span>):
        gating_output, indices = self.router(x)
        final_output = torch.zeros_like(x)

        <span class="hljs-comment"># Reshape inputs for batch processing</span>
        flat_x = x.view(-<span class="hljs-number">1</span>, x.size(-<span class="hljs-number">1</span>))
        flat_gating_output = gating_output.view(-<span class="hljs-number">1</span>, gating_output.size(-<span class="hljs-number">1</span>))

        <span class="hljs-comment"># Process each expert in parallel</span>
        <span class="hljs-keyword">for</span> i, expert <span class="hljs-keyword">in</span> <span class="hljs-built_in">enumerate</span>(self.experts):
            <span class="hljs-comment"># Create a mask for the inputs where the current expert is in top-k</span>
            expert_mask = (indices == i).<span class="hljs-built_in">any</span>(dim=-<span class="hljs-number">1</span>)
            flat_mask = expert_mask.view(-<span class="hljs-number">1</span>)

            <span class="hljs-keyword">if</span> flat_mask.<span class="hljs-built_in">any</span>():
                expert_input = flat_x[flat_mask]
                expert_output = expert(expert_input)

                <span class="hljs-comment"># Extract and apply gating scores</span>
                gating_scores = flat_gating_output[flat_mask, i].unsqueeze(<span class="hljs-number">1</span>)
                weighted_output = expert_output * gating_scores

                <span class="hljs-comment"># Update final output additively by indexing and adding</span>
                final_output[expert_mask] += weighted_output.squeeze(<span class="hljs-number">1</span>)

        <span class="hljs-keyword">return</span> final_output
</code></pre>
<p>This can now be combined with multihead self attention to create a sparse MoE transformer block.</p>
<pre><code class="language-python"><span class="hljs-keyword">class</span> <span class="hljs-title class_">SparseMoEBlock</span>(nn.Module):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, n_embd, num_heads, num_experts, top_k, dropout=<span class="hljs-number">0.1</span>, is_decoder=<span class="hljs-literal">False</span></span>):
        <span class="hljs-built_in">super</span>().__init__()

        <span class="hljs-comment"># Layer normalization for the input to the attention layer</span>
        self.ln1 = nn.LayerNorm(n_embd)

        <span class="hljs-comment"># Multi-head attention module</span>
        self.attn = MultiHeadAttention(n_embd, num_heads, dropout, is_decoder)

        <span class="hljs-comment"># Layer normalization for the input to the FFN</span>
        self.ln2 = nn.LayerNorm(n_embd)

        <span class="hljs-comment"># Feed-forward neural network (FFN)</span>
        self.sparseMoE = SparseMoE(n_embd, num_experts, top_k)

    <span class="hljs-keyword">def</span> <span class="hljs-title function_">forward</span>(<span class="hljs-params">self, x</span>):
        original_x = x  <span class="hljs-comment"># Save the input for the residual connection</span>

        <span class="hljs-comment"># Apply layer normalization to the input</span>
        x = self.ln1(x)

        <span class="hljs-comment"># Apply multi-head attention</span>
        attn_output = self.attn(x)

        <span class="hljs-comment"># Add the residual connection (original input) to the attention output</span>
        x = original_x + attn_output

        <span class="hljs-comment"># Apply layer normalization to the input to the FFN</span>
        x = self.ln2(x)

        <span class="hljs-comment"># Apply the FFN</span>
        sparseMoE_output = self.sparseMoE(x)

        <span class="hljs-comment"># Add the residual connection (input to FFN) to the FFN output</span>
        x = x + sparseMoE_output

        <span class="hljs-keyword">return</span> x
</code></pre>
<p>Now we combine the sparse MoE transformer architecture language decoder model that has been modified to accomodate 'visual tokens' created by the vision-language projector module. Typically, the decoder language model (sparse MoE or dense) will be kept unmodified and will receive embeddings, I've incorportated the vision-langauge projector to the model architecture to keep things simple. A detailed write-up is found in this blog: <a href="https://huggingface.co/blog/AviSoori1x/seemore-vision-language-model">https://huggingface.co/blog/AviSoori1x/seemore-vision-language-model</a></p>
<pre><code class="language-python"><span class="hljs-keyword">class</span> <span class="hljs-title class_">MoEDecoderLanguageModel</span>(nn.Module):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, n_embd, image_embed_dim, vocab_size, num_heads, n_layer, num_experts, top_k, use_images=<span class="hljs-literal">False</span></span>):
        <span class="hljs-built_in">super</span>().__init__()

        self.use_images = use_images

        <span class="hljs-comment"># Token embedding table</span>
        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)

        <span class="hljs-comment"># Position embedding table</span>
        self.position_embedding_table = nn.Embedding(<span class="hljs-number">1000</span>, n_embd)

        <span class="hljs-keyword">if</span> use_images:
            <span class="hljs-comment"># Image projection layer to align image embeddings with text embeddings</span>
            self.image_projection = MultiModalProjector(n_embd, image_embed_dim)

        <span class="hljs-comment"># Stack of transformer decoder blocks</span>
        self.sparseMoEBlocks = nn.Sequential(*[SparseMoEBlock(n_embd, num_heads, num_experts, top_k, is_decoder=<span class="hljs-literal">True</span>) <span class="hljs-keyword">for</span> _ <span class="hljs-keyword">in</span> <span class="hljs-built_in">range</span>(n_layer)])

        <span class="hljs-comment"># Final layer normalization</span>
        self.ln_f = nn.LayerNorm(n_embd)

        <span class="hljs-comment"># Language modeling head</span>
        self.lm_head = nn.Linear(n_embd, vocab_size)

    <span class="hljs-keyword">def</span> <span class="hljs-title function_">forward</span>(<span class="hljs-params">self, idx, image_embeds=<span class="hljs-literal">None</span>, targets=<span class="hljs-literal">None</span></span>):
        <span class="hljs-comment"># Get token embeddings from the input indices</span>
        tok_emb = self.token_embedding_table(idx)

        <span class="hljs-keyword">if</span> self.use_images <span class="hljs-keyword">and</span> image_embeds <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-literal">None</span>:
            <span class="hljs-comment"># Project and concatenate image embeddings with token embeddings</span>
            img_emb = self.image_projection(image_embeds).unsqueeze(<span class="hljs-number">1</span>)
            tok_emb = torch.cat([img_emb, tok_emb], dim=<span class="hljs-number">1</span>)

        <span class="hljs-comment"># Get position embeddings</span>
        pos_emb = self.position_embedding_table(torch.arange(tok_emb.size(<span class="hljs-number">1</span>), device=device)).unsqueeze(<span class="hljs-number">0</span>)

        <span class="hljs-comment"># Add position embeddings to token embeddings</span>
        x = tok_emb + pos_emb

        <span class="hljs-comment"># Pass through the transformer decoder blocks</span>
        x = self.sparseMoEBlocks(x)

        <span class="hljs-comment"># Apply final layer normalization</span>
        x = self.ln_f(x)

        <span class="hljs-comment"># Get the logits from the language modeling head</span>
        logits = self.lm_head(x)

        <span class="hljs-keyword">if</span> targets <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-literal">None</span>:
            <span class="hljs-keyword">if</span> self.use_images <span class="hljs-keyword">and</span> image_embeds <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-literal">None</span>:
                <span class="hljs-comment"># Prepare targets by concatenating a dummy target for the image embedding</span>
                batch_size = idx.size(<span class="hljs-number">0</span>)
                targets = torch.cat([torch.full((batch_size, <span class="hljs-number">1</span>), -<span class="hljs-number">100</span>, dtype=torch.long, device=device), targets], dim=<span class="hljs-number">1</span>)

            <span class="hljs-comment"># Compute the cross-entropy loss</span>
            loss = F.cross_entropy(logits.view(-<span class="hljs-number">1</span>, logits.size(-<span class="hljs-number">1</span>)), targets.view(-<span class="hljs-number">1</span>), ignore_index=-<span class="hljs-number">100</span>)
            <span class="hljs-keyword">return</span> logits, loss

        <span class="hljs-keyword">return</span> logits

    <span class="hljs-keyword">def</span> <span class="hljs-title function_">generate</span>(<span class="hljs-params">self, idx, image_embeds, max_new_tokens</span>):
        <span class="hljs-comment"># The autoregressive character level generation function is just like in any other decoder model implementation</span>

        <span class="hljs-keyword">return</span> generated
</code></pre>
<p>Now that we have our three key components, we can put it all together into a sparse Mixture of Experts Vision Language Model. The full implementation is given below. If you were to remove the assert statements for error handling, this is very simple. Coming back full circle to the outline I’ve given at the beginning of the blog, all that’s happening here is:</p>
<ol>
<li><p>Get image features from the vision encoder (Here it’s a vision transformer, but it could be any model that could generate features from an image input such as a ResNet or a traditional convolutional neural network (needless to say performance may suffer))</p>
</li>
<li><p>A projection module for projecting image tokens to the same embedding space as text embeddings for the decoder (this projector is integrated with the decoder in this implementation)</p>
</li>
<li><p>A decoder language model with a sparseMoE architecture for generating text conditioned on a preceding image.</p>
</li>
</ol>
<pre><code class="language-python"><span class="hljs-keyword">class</span> <span class="hljs-title class_">VisionMoELanguageModel</span>(nn.Module):
    <span class="hljs-keyword">def</span> <span class="hljs-title function_">__init__</span>(<span class="hljs-params">self, n_embd, image_embed_dim, vocab_size, n_layer, img_size, patch_size, num_heads, num_blks, emb_dropout, blk_dropout, num_experts, top_k</span>):
        <span class="hljs-built_in">super</span>().__init__()

        <span class="hljs-comment"># Set num_hiddens equal to image_embed_dim</span>
        num_hiddens = image_embed_dim

        <span class="hljs-comment"># Assert that num_hiddens is divisible by num_heads</span>
        <span class="hljs-keyword">assert</span> num_hiddens % num_heads == <span class="hljs-number">0</span>, <span class="hljs-string">"num_hiddens must be divisible by num_heads"</span>

        <span class="hljs-comment"># Initialize the vision encoder (ViT)</span>
        self.vision_encoder = ViT(img_size, patch_size, num_hiddens, num_heads, num_blks, emb_dropout, blk_dropout)

        <span class="hljs-comment"># Initialize the language model decoder (DecoderLanguageModel)</span>
        self.decoder = MoEDecoderLanguageModel(n_embd, image_embed_dim, vocab_size, num_heads, n_layer,num_experts, top_k, use_images=<span class="hljs-literal">True</span>)

    <span class="hljs-keyword">def</span> <span class="hljs-title function_">forward</span>(<span class="hljs-params">self, img_array, idx, targets=<span class="hljs-literal">None</span></span>):
        <span class="hljs-comment"># Get the image embeddings from the vision encoder</span>
        image_embeds = self.vision_encoder(img_array)

        <span class="hljs-comment"># Check if the image embeddings are valid</span>
        <span class="hljs-keyword">if</span> image_embeds.nelement() == <span class="hljs-number">0</span> <span class="hljs-keyword">or</span> image_embeds.shape[<span class="hljs-number">1</span>] == <span class="hljs-number">0</span>:
            <span class="hljs-keyword">raise</span> ValueError(<span class="hljs-string">"Something is wrong with the ViT model. It's returning an empty tensor or the embedding dimension is empty."</span>)

        <span class="hljs-keyword">if</span> targets <span class="hljs-keyword">is</span> <span class="hljs-keyword">not</span> <span class="hljs-literal">None</span>:
            <span class="hljs-comment"># If targets are provided, compute the logits and loss</span>
            logits, loss = self.decoder(idx, image_embeds, targets)
            <span class="hljs-keyword">return</span> logits, loss
        <span class="hljs-keyword">else</span>:
            <span class="hljs-comment"># If targets are not provided, compute only the logits</span>
            logits = self.decoder(idx, image_embeds)
            <span class="hljs-keyword">return</span> logits

    <span class="hljs-keyword">def</span> <span class="hljs-title function_">generate</span>(<span class="hljs-params">self, img_array, idx, max_new_tokens</span>):
        <span class="hljs-comment"># Get the image embeddings from the vision encoder</span>
        image_embeds = self.vision_encoder(img_array)

        <span class="hljs-comment"># Check if the image embeddings are valid</span>
        <span class="hljs-keyword">if</span> image_embeds.nelement() == <span class="hljs-number">0</span> <span class="hljs-keyword">or</span> image_embeds.shape[<span class="hljs-number">1</span>] == <span class="hljs-number">0</span>:
            <span class="hljs-keyword">raise</span> ValueError(<span class="hljs-string">"Something is wrong with the ViT model. It's returning an empty tensor or the embedding dimension is empty."</span>)

        <span class="hljs-comment"># Generate new tokens using the language model decoder</span>
        generated_tokens = self.decoder.generate(idx, image_embeds, max_new_tokens)
        <span class="hljs-keyword">return</span> generated_tokens
</code></pre>
<p>Now back to where we started. The above VisionMoELanguageModel class neatly wraps up all the components we set out to put together.</p>
<img alt="seemore" height="500" src="https://github.com/AviSoori1x/seemore/blob/main/images/moevlm.png?raw=true" width="500"/>
<p>  The training loop is exactly as you see in the original vision language blog linked at the top. Please check out seeMoE_from_Scratch.ipynb in the following repo: <a href="https://github.com/AviSoori1x/seemore" rel="nofollow">https://github.com/AviSoori1x/seemore</a></p>
<p>PS: There are newer approaches such as mixed-modal early-fusion models e.g. <a href="https://arxiv.org/abs/2405.09818" rel="nofollow">https://arxiv.org/abs/2405.09818</a>. I plan to implement a simplistic version of this in the future. </p>
<p>Thanks for reading!</p>
