I Built a Multimodal Embedding Model From Scratch on an RTX 4060 (Text, Image, Audio, and Video…

I Built a Multimodal Embedding Model From Scratch on an RTX 4060 (Text, Image, Audio, and Video, One Vector Space, 8GB VRAM)

How I reproduced the architecture behind Jina AI’s GELATO on a single consumer laptop GPU, and what seven rounds of failed experiments taught me about where multimodal embedding models actually break

I trained a text, image, audio, and video embedding model that shares one 768 dimensional vector space, on a single RTX 4060 laptop GPU with 8GB of VRAM. Peak memory usage never exceeded 2.7GB. No cloud compute, no rented A100s, no enterprise cluster.

This post is a deep dive into reproducing GELATO, the architecture behind Jina AI’s jina-embeddings-v5-omni model. Huge credit to the Jina AI team for the original research. Their paper on arXiv describes how to build a multimodal embedding model while training less than half a percent of the parameter count, and this project is my attempt to actually implement it end to end on hardware anyone reading this probably already owns.

If you are searching for how to build a multimodal embedding model, how to train an embedding model on a consumer GPU, or whether frozen encoder projector training actually works in practice, this is a full walkthrough with real numbers, real failures, and the actual code.

I called the finished project QuadEmbed, after the four modalities it covers. Weights are on Hugging Face, code is on GitHub.

The core idea: contrastive alignment with frozen towers

Most multimodal training approaches update the entire network. GELATO does the opposite, and the reasoning is worth understanding before looking at any code.

A pretrained text encoder already has a well structured embedding space. Semantically similar sentences are already close together, dissimilar ones already far apart. That geometry took enormous compute to produce. If you fine tune that encoder while simultaneously trying to align a vision encoder to it, you are optimizing against a moving target, and you risk destroying the geometry you actually wanted to keep. The paper’s own ablation confirms this: unfreezing the encoders performs worse than leaving them alone.

So GELATO freezes three pretrained encoders completely and trains only small projection heads that map vision and audio features into the text encoder’s existing space. The text space is the anchor. Everything else learns to point at it.

Here is the architecture as I implemented it:

The exact models:

• Text (frozen anchor): jinaai/jina-embeddings-v5-text-nano, 239M parameters, outputs 768 dimensions
• Vision (frozen): google/siglip2-base-patch16-naflex, outputs 768 dimensions per patch
• Audio (frozen): openai/whisper-large-v3, encoder half only, outputs 1280 dimensions per frame

Trainable surface area: a vision projector at 2.36M parameters and an audio projector at 0.98M. Roughly 3.3M trainable against nearly a billion frozen.

One detail worth flagging: the dimensions line up with the paper’s stated numbers without any adjustment. SigLIP2 patches are 768 dimensional, and merging four of them gives exactly the 3072 input width the paper specifies for the nano vision projector. Whisper’s encoder outputs exactly 1280, matching the paper’s audio projector input. That was a strong early signal the substitution was architecturally faithful rather than merely close.

The loss function: bidirectional InfoNCE plus Matryoshka

The training objective is contrastive. Given a batch of matched image-caption pairs, the model should score the correct pairing higher than all the mismatched pairings in that same batch. Every other item in the batch acts as a negative example, which is why batch size matters so much for contrastive training.

Bidirectional means the loss is computed in both directions and averaged: text retrieving image, and image retrieving text. Here is the actual implementation:

def bidirectional_infonce(a, b, temperature=0.02):
    """a, b: [B, D] paired embeddings, same batch order = positives."""
    a = F.normalize(a, dim=-1)
    b = F.normalize(b, dim=-1)
    logits = a @ b.T / temperature          # [B, B] similarity matrix
    labels = torch.arange(a.shape[0], device=a.device)  # diagonal = correct pairs
    loss_a2b = F.cross_entropy(logits, labels)
    loss_b2a = F.cross_entropy(logits.T, labels)
    return (loss_a2b + loss_b2a) / 2

The trick is that the correct pairings sit exactly on the diagonal of the similarity matrix, so the labels are just arange(batch_size). Cross entropy over rows gives you text to image, over columns gives you image to text.

Layered on top is Matryoshka representation learning, which trains the embedding so that truncated prefixes remain usable embeddings on their own. If you only need 64 dimensions for a fast approximate search index, you slice the first 64 and it still works:

NANO_MATRYOSHKA_DIMS = (32, 64, 128, 256, 768)

def matryoshka_infonce(a, b, dims=NANO_MATRYOSHKA_DIMS, temperature=0.02):
    losses = []
    for d in dims:
        if d > a.shape[-1]:
            continue
        losses.append(bidirectional_infonce(a[..., :d], b[..., :d], temperature))
    return torch.stack(losses).mean()

Temperature is 0.02, matching the paper. Lower temperature sharpens the softmax, penalizing near misses harder.

The training step, and why it fits in 8GB

Here is the actual training loop body:

with torch.no_grad():
    patch_tokens, mask, spatial_shapes = vision_encoder.patch_tokens(images)
    text_embeds = text_encoder.embed(captions)

image_embeds = vision_projector(patch_tokens.float(), mask, spatial_shapes)
loss = matryoshka_infonce(text_embeds.float(), image_embeds)

optimizer.zero_grad()
loss.backward()
optimizer.step()

Look at what torch.no_grad() covers. Both frozen encoders run their forward passes inside it, so PyTorch never builds an autograd graph for them, never stores intermediate activations for backprop, and never allocates optimizer state for their parameters. The backward pass only traverses the projector.

That is the entire memory trick. Activation storage during backprop, not weight storage, is what usually blows up VRAM during training. Eliminating it for 99.6% of the network is why this runs comfortably on a laptop:

Video for free: no video encoder exists

There is no video model anywhere in this system. A video is sampled into frames, each frame goes through the vision encoder and the already trained vision projector, and the resulting vectors are mean pooled over time:

@torch.no_grad()
def embed_video(vision_encoder, vision_projector, video_path, device, num_frames=4):
    frames = sample_frames(video_path, num_frames=num_frames)
    patch_tokens, mask, spatial_shapes = vision_encoder.patch_tokens(frames)
    frame_embeds = vision_projector(patch_tokens.float(), mask, spatial_shapes)
    return frame_embeds.mean(dim=0)

That is the whole video pipeline. It works because the projector already learned to map visual features into text space, and a video frame is just an image. My zero-training video baseline hit 26% R@1 against a 2% chance baseline before I trained on a single video.

Datasets

Everything is public, pulled from the Hugging Face Hub:

Final vision training set: about 172,000 image-caption pairs across two distinct visual domains, curated photography and web imagery.

Held-out evaluation always uses the Flickr8k test split, regardless of what was added to training, so every number across all seven rounds is directly comparable.

Seven rounds of controlled experiments

Each round changed one variable. Here is what happened.

Rounds 1 and 2, the baseline. The architecture works. All three modalities retrieved far above chance. Scaling data and steps pushed audio to 67/97/100 R@1/R@5/R@10, and video from 26% to 40% R@1 on held-out clips.

Round 3, a real bug in the patch merge. Vision R@1 sat at 13% across two rounds. Investigating, I found the projector’s 2x2 spatial merge was not merging spatially adjacent patches at all. SigLIP2’s NaFlex mode packs variable resolution images, so the flattened token sequence does not have a fixed row width, and grouping four consecutive tokens can straddle unrelated rows. The correct implementation needs the per image grid shape, which the processor returns as spatial_shapes but which I was never threading through:

def _spatial_merge_one(self, tokens, h, w):
    # tokens: [N, D] for one image; first h*w entries are the valid grid
    d = tokens.shape[-1]
    h2, w2 = h - (h % 2), w - (w % 2)   # drop odd trailing row/col
    grid = tokens[: h * w].view(h, w, d)[:h2, :w2]
    blocks = grid.reshape(h2 // 2, 2, w2 // 2, 2, d)
    return blocks.permute(0, 2, 1, 3, 4).reshape(-1, 4 * d)

I fixed it, retrained from scratch since the old checkpoints were now semantically incompatible, and R@1 did not move. A genuine correctness bug, fixed properly, with zero measurable effect on retrieval quality.

Round 4, batch size and steps. Batch 64 to 128, steps 2500 to 6000. Training loss dropped from 0.61 to 0.56. Held-out recall got slightly worse. Lower training loss with flat or declining validation metrics is the textbook overfitting signature.

Round 5, domain diversity and a NaN incident. I added Conceptual Captions for genuinely different imagery. CC3M only ships URLs, not image bytes, so I wrote a concurrent downloader: 32 worker threads, 6 second timeout, skip on failure. 30,000 attempts yielded 17,293 images, a 57.6% hit rate.

The first training run on that data went to nan around step 850 and stayed there. The cause was two images out of 17,293 that were essentially one pixel wide, tracking pixels served instead of real content, which still decoded as valid JPEGs and passed a naive byte-length check. Feeding a 1250x1 image into NaFlex's patch grid math produces NaN, and since the projector updates every step, that NaN propagated into the weights and silently destroyed the rest of the run. The fix is two lines in the downloader:

w, h = img.size
if min(w, h) < 32 or max(w, h) / min(w, h) > 6:
    return idx, caption, False   # degenerate image, skip it

After filtering and retraining, R@5 and the reverse retrieval direction hit new highs. R@1 stayed at 13.0%.

Round 6, projector capacity. If a single linear layer was the constraint, more capacity should help. I swapped it for a 2 layer MLP, 3072 to 1536 to 768 with GELU. Training loss collapsed to 0.13, by far the best fit of any round. Held-out recall got worse on every metric. More capacity on limited data buys memorization, not generalization. This also independently confirms the paper’s ablation finding that constrained architectures win when data is the bottleneck.

Round 7, just more data. COCO’s train split has 182 parquet shards and I had been using 2 to 4. Scaling to 40 shards brought the training set to ~172k pairs, a 3.2x increase, while deliberately reducing epoch count from ~14 to ~7 to avoid repeating round 4’s mistake. Every metric moved at once.

Reading the results honestly

Final numbers: R@1 13.7%, R@5 68.6%, R@10 81.1%, against a 0.1% random chance baseline on 1024 candidates.

The R@1 improvement in round 7 deserves scrutiny. With n=1024, the standard error on a proportion near 13% is roughly one percentage point. So 13.3% to 13.7% is, in isolation, inside the noise band. I am not going to claim a clean win from that number alone.

What makes round 7 credible is the correlation structure. Every metric moved the same direction simultaneously, including both retrieval directions. In rounds 4 and 6, metrics moved in different directions while training loss improved, which is what overfitting looks like. Round 7 shows the opposite pattern. Coordinated movement across independent metrics is much harder to produce by chance than a single metric drifting.

What this actually proves

Five interventions did nothing to R@1: fixing a real architecture bug, doubling the negative pool, adding a genuinely different visual domain, and increasing projector capacity. One intervention moved it: more training data.

For anyone building on frozen encoder architectures, the practical implication is that your model capacity is almost certainly not your problem. The projector is a coordinate transform between two spaces that are already well formed. Learning a coordinate transform does not require much capacity, but it does require enough examples to constrain the transform properly. At 40k pairs I was underdetermined. At 172k pairs I was less so. The paper trains on orders of magnitude more than that, which is most of the remaining gap.

The corollary is that debugging effort should be spent proportionally. I spent an entire round rebuilding a projector that turned out to be innocent, and one round adding data that turned out to be the answer.

It is on PyPI, so you can go from nothing to embeddings in two commands:

pip install quadembed          # text + image + audio
pip install quadembed[video]   # adds video support

from quadembed import QuadEmbed
from PIL import Image

model = QuadEmbed.from_pretrained()   # downloads the weights on first run

text_embeds = model.embed_text(["a dog running on the beach"])
image_embeds = model.embed_image([Image.open("photo.jpg").convert("RGB")])
similarity = text_embeds @ image_embeds.T   # L2-normalized, so this is cosine

Each frozen encoder costs memory and download time, so load only the ones you need:

model = QuadEmbed.from_pretrained(modalities=("text", "vision"))   # skip audio

Audio takes mono float32 arrays at 16 kHz via embed_audio(), and video takes a file path via embed_video_file().

One caveat worth stating plainly: the frozen text encoder is licensed CC-BY-NC, and that carries through to this package and its weights. Research and educational use is fine, commercial use is not without a separate license from Jina AI.

• Package: pypi.org/project/quadembed
• Model weights: Mithil-AI/quadembed-nano on Hugging Face
• Full source, training scripts, and the complete round by round experiment log: mithilai/QuadEmbed on GitHub

If you build something similar, or hit your own version of the stuck metric problem, I would genuinely like to hear about it.