💔The SwiftFox architecture has gone through SO many iterations and existential crises (: and no, today I'm not here to debate names 😭🙏.

Tbh, it was lowkey such a fun and crazy project to build, and I'm definitely going to keep developing it. If you remember my October 28 post, I attached an image to my earlier prototype, asked it what was inside, and it literally hallucinated complete non-sense code with a perplexity of 850.309159... like, what else could I even expect?! 💔 That kind of pain makes you want to rewrite everything from absolute scratch.

So I sat down and built SwiftFox: a multimodal language model implemented completely in PyTorch that ingests both image patches and raw text tokens and autoregressively generates clean responses. The architecture is built to work seamlessly whether you give it pure Text → Text or multiple visual inputs for Image + Text → Text.

And before anyone asks: YES, it already has native multi-image cross-frame reasoning built right into its core! In my training pipeline, I pass up to 3 images at once (and structurally it can scale to even more). It doesn't process images in isolation — it flattens patches from ALL images into the exact same self-attention sequence, meaning the model can literally compare frames, track visual changes, and solve tasks like Spot-The-Diff natively 😭🙏.

call:SwiftFox flattens all patches across multiple images and text tokens into a single autoregressive sequence, calculating dense cross-image attention with RoPE, and routing tokens through an Adaptive Entropy Mixture of Experts (MoE) powered by SwiGLU.

SwiftFox Pipeline & Multimodal Dataflow
From Raw Inputs to Autoregressive Token Generation
01 / INPUTS

Dual-Modal Ingestion

Supports pure Text → Text or Image + Text → Text (up to 3 images per context).

IMG_1..3+Text Prompt
02 / EMBEDDINGS

Residual 2D + nn.Embedding

Images are sliced into patches and projected via 2D Residual Blocks (1x1 Conv). Text is mapped through nn.Embedding.

[IMG_1]..[IMG_N] [TXT_1]..[TXT_M]
03 / ATTENTION

RoPE Multi-Head Attention

Rotary Position Embeddings rotate Q & K pairs. Invariant to absolute position, ensuring natural visual-textual relative decay.

RoPESelf-AttentionCausal Mask
04 / ROUTING & MoE

Adaptive Top-K with SwiGLU

Shannon entropy H(x) dynamically adjusts K experts. Next-token prediction trained with L_CE + L_LB.

SwiGLUEntropy RouterLoad Balance

1. The Real Timeline (:.

School Day1.5 HOURS

3 Random CS Classes

Was sitting through 3 computer science classes at school, got bored, and hammered out the core multimodal architecture in an hour and a half.

Evening+1 HOUR

1 Hour at Home

Got back home, spent about one more hour polishing the patching pipeline and setting up the training loop.

Next 2 Days2 DAYS

The 48h Training Run

Trained it for 2 straight days. Saw that it was actually learning and converging, but realized it could be way better.

~2 Weeks LaterMOE

Why Not MoE?

Came back to the code after about a week and a half or two. Decided to slap an Adaptive Mixture of Experts with SwiGLU onto it because honestly... why not (: that's literally it.

2. Multi-Image Unified Sequence: How Cross-Frame Reasoning Actually Works.

Let's clear this up right away because this is one of the coolest parts of the whole model: how does SwiftFox look at multiple images and understand relationships between them? (: In my notebook (very_very_random_file.ipynb), I mapped out the exact sequence design:

sequence_layout.txt
text
Text chat mode:
[<user>] {input} [<assistant>] {generated_text}

Multi-image mode (up to 3 images, expandable to more):
[IMG1_1] ... [IMG1_16]  [IMG2_1] ... [IMG2_16]  [<user>] [T1] ... [TN]  [<assistant>] [G1] ... [GM]

Think about what this means for cross-frame reasoning: most architectures force you to run separate image encoders, pool each image into a single vector, or use clunky cross-attention fusion adapters that lose spatial granularity. In SwiftFox, every image is sliced into 16 distinct patches (64x64). If you pass 2 images, you get 32 visual patch tokens. If you pass 3 images, you get 48 visual patch tokens.

Because all 48 visual tokens are concatenated alongside the text prompt into the same Multi-Head Attention block, the transformer learns relationships between image patches *across different frames*! Image 2's patches can directly query Image 1's patches. That's how SwiftFox tackles datasets like Spot-the-diff and multi-image instructions: the model literally sees what changed from frame A to frame B.

Here is the exact visual encoder implementation from my notebook handling 5D tensors (B, num_imgs, C, H, W) and slicing them into 16 spatial patches:

image_encoder.py
python
import torch
import torch.nn as nn

class ResidualBlock2D(nn.Module):
    def __init__(self, in_, out_, stride_=1):
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_, out_, kernel_size=3, stride=stride_, padding=1),
            nn.BatchNorm2d(out_),
            nn.ReLU(),
            nn.Conv2d(in_, out_, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(out_)
        )
        self.identity_ = nn.Identity()
        if in_ != out_ or stride_ != 1:
            self.identity_ = nn.Sequential(
                nn.Conv2d(in_, out_, kernel_size=1, stride=stride_),
                nn.BatchNorm2d(out_)
            )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.conv1(x) + self.identity_(x))

class ImageEncoder(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.more_layers = nn.Sequential(
            ResidualBlock2D(3, 64),
            ResidualBlock2D(64, 128),
            ResidualBlock2D(128, 256),
            nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.AdaptiveMaxPool2d((8, 8)),
            nn.Flatten(),
            nn.Linear(256 * 8 * 8, d_model)
        )

    def forward(self, x, batch_size=None):
        if x.dim() == 5:
            batch_size = x.size(0)
            x = x.view(-1, x.size(2), x.size(3), x.size(4))
        if batch_size is None:
            batch_size = x.size(0)

        total_images = x.size(0)
        H_patch = x.size(2) // 4
        W_patch = x.size(3) // 4
        x = x.reshape(total_images, 16, x.size(1), H_patch, W_patch)
        x = x.reshape(16 * total_images, x.size(1), H_patch, W_patch)
        x = self.more_layers(x)
        
        return x.reshape(batch_size, -1, x.size(-1))

3. Rotary Position Embeddings (RoPE) & Attention.

Standard absolute positional encodings get super awkward when you glue multiple images in front of text tokens. Static 1D sinusoidal position vectors have no clue how to handle spatial patches, let alone multiple sequential images.

Instead, I implemented RoPE (Rotary Position Embeddings) directly on Query and Key vectors. RoPE rotates pairs of coordinates in feature space using complex exponentials. When Query and Key compute dot-product attention, their absolute positions cancel out, leaving only the relative distance (m - n)! This makes relative spatial proximity between patches and relative causal distance between tokens feel completely natural.

Here is the exact RoPE and Multi-Head Attention module from Cell 3 of very_very_random_file.ipynb:

rope_mha.py
python
import math
import torch
import torch.nn as nn

class RoPE(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        inv_freq = torch.exp(
            torch.arange(0, head_dim, 2) * (-1) * math.log(10000) / head_dim
        )
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, x):
        B, T, H, D = x.shape
        position = torch.arange(0, T, device=x.device)
        freq = torch.outer(position, self.inv_freq)

        cos_ = freq.cos().unsqueeze(0).unsqueeze(0)
        sin_ = freq.sin().unsqueeze(0).unsqueeze(0)

        x_even = x[..., ::2]
        x_odd = x[..., 1::2]
        x_rot = x_even * cos_ - x_odd * sin_
        y_rot = x_even * sin_ + x_odd * cos_
        return torch.stack((x_rot, y_rot), dim=-1).flatten(-2)

class MHA(nn.Module):
    def __init__(self, d_model, nhead, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.head_dim = d_model // nhead

        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        self.wo = nn.Linear(d_model, d_model)
        self.d1 = nn.Dropout(dropout)
        self.d2 = nn.Dropout(dropout)
        self.rope = RoPE(self.head_dim)

    def forward(self, x, src_key_padding_mask=None, src_mask=None):
        B, T, F = x.shape
        k = self.wk(x).transpose(1, 2).view(B, T, self.nhead, self.head_dim)
        q = self.wq(x).transpose(1, 2).view(B, T, self.nhead, self.head_dim)
        v = self.wv(x).transpose(1, 2).view(B, T, self.nhead, self.head_dim)

        k = self.rope(k)
        q = self.rope(q)

        score = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if src_key_padding_mask is not None:
            src_key_padding_mask = src_key_padding_mask.unsqueeze(1).unsqueeze(2)
            score = score.masked_fill(src_key_padding_mask == float("-inf"), float("-inf"))
        if src_mask is not None:
            score = score.masked_fill(src_mask, float("-inf"))

        score = torch.softmax(score, dim=-1)
        score = self.d1(score)
        score = score.view(B, T, F).contiguous()
        return self.d2(self.wo(score))

4. SwiGLU & The Adaptive Entropy MoE.

Now we get to the real meat of the architecture: the Mixture of Experts (MoE) backbone. Why run one monolithic feedforward layer when you can specialize computation across 64 specialized experts?

SwiGLU: Gated Feedforward

First off, vanilla ReLU or GELU feedforward networks are cool, but modern LLMs get way better gradient flow and expressive capacity out of SwiGLU. The SwiGLU block splits the projection into a gated pair:

swiglu.py
python
import torch.nn as nn
import torch.nn.functional as F

class SwiGLU(nn.Module):
    def __init__(self, d_model, dim_feedforward):
        super().__init__()
        self.w1 = nn.Linear(d_model, dim_feedforward)
        self.w2 = nn.Linear(d_model, dim_feedforward)
        self.w3 = nn.Linear(dim_feedforward, d_model)

    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))

The Adaptive Top-K Router

Here was my biggest realization: in standard MoE (like Mixtral 8x7B), K is strictly fixed to e.g. top-2. But why should every single token wake up the exact same number of experts?! 😭🙏

If a token is dead obvious (like a period or a common preposition), the router distribution has high confidence, meaning Shannon entropy H(x) is almost 0. Waking up 4 or 8 heavy experts is just wasting GPU cycles! But when the token is tricky, ambiguous, or bridging visual and text features, the distribution flattens, entropy spikes, and you want to call in a whole squad of experts.

So I formulated the Adaptive Top-K rule:

Adaptive Top-K Dynamic Allocation
K_{\text{adaptive}} = K_{\text{min}} + \lfloor H(x) \cdot (K_{\text{max}} - K_{\text{min}}) \rfloor

Scales expert capacity dynamically based on normalized Shannon routing entropy.

K_{\text{min}}:Base floor of active experts (e.g. 4)
K_{\text{max}}:Ceiling of active experts (e.g. 64)
H(x):Normalized Shannon entropy: -sum(p * ln(p + 1e-9)) / ln(K_{\text{max}})

And here is the actual PyTorch implementation from Cell 10 of very_very_random_file.ipynb, routing experts dynamically and accumulating outputs with output.index_put_:

adaptive_moe.py
python
class MoE(nn.Module):
    def __init__(self, input_dim, hidden_dim, min_k=4, max_k=64):
        super().__init__()
        assert min_k < max_k
        self.router = nn.Linear(input_dim, max_k)
        self.experts = nn.ModuleList([
            SwiGLU(input_dim, hidden_dim) for _ in range(max_k)
        ])
        self.max_k = max_k
        self.min_k = min_k

    def forward(self, x):
        B, T, F = x.shape
        logits = self.router(x)
        probs = torch.softmax(logits, dim=-1)
        
        entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1) / math.log(self.max_k)
        adaptive_k = self.min_k + torch.floor(entropy * (self.max_k - self.min_k)).long()

        top_probs, top_indices = torch.topk(probs, adaptive_k, dim=-1)
        weights = top_probs / top_probs.sum(dim=-1, keepdim=True)

        output = torch.zeros(B, T, F, device=x.device, dtype=x.dtype)

        for i_exp, expert in enumerate(self.experts):
            b_indices, t_indices, k_indices = torch.where(top_indices == i_exp)
            if b_indices.numel() == 0:
                continue
            expert_weights = weights[b_indices, t_indices, k_indices]
            expert_input = probs[b_indices, t_indices]
            expert_output = expert(expert_input)

            output.index_put_(
                (b_indices, t_indices),
                expert_weights * expert_output,
                accumulate=True
            )

        return x + output, probs, top_indices

5. Loss: Cross-Entropy + Load Balancing.

Anyone who has trained MoEs knows the infamous nightmare: router collapse 💔.

If you don't incentivize the router to spread the work, it will quickly get lazy and send 99% of all tokens to Expert 0, while the remaining 63 experts sit around with zero gradients. To prevent this, I added an auxiliary Load Balancing Loss alongside the autoregressive cross-entropy loss:

Total SwiftFox Optimization Objective
\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{CE}} + \mathcal{L}_{\text{LB}}

Joint objective combining autoregressive token prediction with router distribution balance.

\mathcal{L}_{\text{CE}}:Standard Cross Entropy loss computed on targets text[1:]
\mathcal{L}_{\text{LB}}:Load Balance Loss: E * sum(P_i * f_i)
P_i:Average routing probability across all tokens for expert i
f_i:Actual fraction of tokens assigned to expert i

Here is the auxiliary load balance loss function directly from the notebook:

load_balance.py
python
def load_balance_loss(probs, top_indices, n_e):
    B, T, F = probs.shape
    importance = torch.mean(probs, dim=0)
    count = torch.bincount(top_indices.reshape(-1), minlength=n_e)
    frequency = count.float() / (B * T * F)
    return n_e * torch.sum(importance * frequency)

6. Autoregressive Prediction: Predicting text[1:].

During training, the shifted sequence mechanics work cleanly:

When I feed model(text, images), what I pass into the network is the sequence without the last token. The model is trained to predict text[1:] (the tokens shifted by one step to the right). The loss is calculated exclusively on the generated assistant tokens, conditioning on both the preceding multi-image visual patches and the user instructions.

training_step.py
python
logits, all_probs, all_indices = model(text=input_ids[:, :-1], images=images)

targets = input_ids[:, 1:]
loss_ce = criterion(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))

loss_lb = 0.0
for probs, indices in zip(all_probs, all_indices):
    loss_lb += load_balance_loss(probs, indices, n_e=64)

total_loss = loss_ce + loss_lb
total_loss.backward()
optimizer.step()
scheduler.step()

💔Hope this is not a fail, whatever (: Watching the perplexity drop from that horrific 850.3 to crisp, coherent predictions on both multi-image inputs and text prompts makes the sleepless nights completely worth it.

Having multi-image cross-frame reasoning natively working without heavy external vision backbones is super exciting. Next up: scaling up the parameter count, pushing the context window even further, and releasing the model checkpoints once fine-tuning wraps up. Stay tuned (: !