quadembed-nano / README.md
Mithil-AI's picture
Fix pipeline tag: this is an embedding/retrieval model, not text-to-image generation
0217aee verified
|
Raw
History Blame Contribute Delete
8.22 kB
---
license: cc-by-nc-4.0
pipeline_tag: feature-extraction
tags:
- multimodal
- embeddings
- feature-extraction
- sentence-similarity
- image-text-retrieval
- audio
- video
- retrieval
- contrastive-learning
- gelato
- jina-embeddings
library_name: pytorch
---
# QuadEmbed
A local, from-scratch-trained multimodal embedding model. "Quad" for the
four modalities it covers: text, image, audio, and video, all mapped into
one shared 768-dimensional embedding space, trained end-to-end on a single
8GB consumer GPU (RTX 4060 laptop).
The architecture reproduces **GELATO** (Geometry-preserving Embeddings via
Locked Aligned TOwers), described by Jina AI in the paper behind
[jina-embeddings-v5-omni](https://arxiv.org/abs/2605.08384). Full credit for
the architecture and training recipe goes to that paper; this repo is an
independent, from-scratch reproduction of it, not a copy of Jina's released
weights.
This is not a copy of Jina's released checkpoint β€” the three encoders below
are frozen, publicly available source models, and only two small projector
heads (a few million parameters total) were trained from scratch on public
datasets to align them into a shared space.
πŸ“¦ **`pip install quadembed`** ([PyPI](https://pypi.org/project/quadembed/)) Β·
πŸ’» [Source on GitHub](https://github.com/mithilai/QuadEmbed) Β·
πŸ“ [Full write-up on Medium](https://medium.com/@mithilmaske/i-built-a-multimodal-embedding-model-from-scratch-on-an-rtx-4060-text-image-audio-and-video-ab1fef04f1cd)
(how it was built, and seven rounds of what did and didn't work)
## Architecture
Three frozen encoders, two trained projectors β€” everything large stays
untouched; only the projectors and two small delimiter vectors have
gradients:
| Role | Model (frozen) | Params | Role in this repo |
|---|---|---|---|
| Text (anchor) | [`jinaai/jina-embeddings-v5-text-nano`](https://huggingface.co/jinaai/jina-embeddings-v5-text-nano) | 239M | Defines the target embedding space; never modified |
| Vision | [`google/siglip2-base-patch16-naflex`](https://huggingface.co/google/siglip2-base-patch16-naflex) | ~93M | Patch features, variable resolution (NaFlex) |
| Audio | [`openai/whisper-large-v3`](https://huggingface.co/openai/whisper-large-v3) (encoder only) | ~635M | Frame-level audio features |
| **Vision projector** (trained) | `checkpoints/vision_projector.pt` | 2.36M | LayerNorm → real 2×2 spatial patch merge → linear (3072→768) |
| **Audio projector** (trained) | `checkpoints/audio_projector.pt` | 0.98M | Linear (1280→768) |
**Video has no dedicated encoder or projector.** A video is 4 sampled frames
run through the vision encoder + vision projector, then mean-pooled over
time β€” exactly how GELATO's own paper handles video. `checkpoints/video_projector.pt`
is the vision projector after additional fine-tuning on real video frames
(see caveat below).
Why these particular substitutions: GELATO's paper describes its vision and
audio towers as *adapted from* SigLIP2 and Whisper-large-v3, but doesn't
release those adapted checkpoints publicly β€” so this reproduction uses the
public source models directly. The text encoder is an exact match; Jina
describes v5-omni's text tower as bit-identical to the standalone
`jina-embeddings-v5-text` release.
## Results (measured on this hardware, not estimated)
Cross-modal retrieval recall@k, text-query direction, against a held-out
split (higher is better):
| Modality | R@1 | R@5 | R@10 | n |
|---|---|---|---|---|
| Image | 13.7% | 68.6% | 81.1% | 1024 |
| Audio | 67% | 97% | 100% | 33 |
| Video (held-out) | 40% | 86% | 94% | 50 |
Random-chance R@1 on the 1024-candidate image eval is ~0.1%; all three
modalities land far above chance. Peak VRAM across every training run
stayed under 2.7GB β€” well inside an 8GB budget.
These numbers come after seven rounds of iterating on the vision projector
specifically (architecture fixes, data scale, data diversity, projector
capacity). Short version: architecture wasn't the bottleneck after round 3;
raw training data volume was β€” R@1 only moved once the image-caption
training set was scaled to ~172k pairs (round 7). The full round-by-round
write-up, with what did and didn't move the needle, lives in the source
repo's README and feasibility notes:
[github.com/mithilai/QuadEmbed](https://github.com/mithilai/QuadEmbed), and
as a full narrative write-up on Medium:
[I Built a Multimodal Embedding Model From Scratch on an RTX 4060](https://medium.com/@mithilmaske/i-built-a-multimodal-embedding-model-from-scratch-on-an-rtx-4060-text-image-audio-and-video-ab1fef04f1cd).
## Usage
```bash
pip install quadembed # text + image + audio
pip install quadembed[video] # adds video support
```
```python
from quadembed import QuadEmbed
from PIL import Image
model = QuadEmbed.from_pretrained() # downloads these weights automatically
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 # already L2-normalized -> cosine similarity
```
Each encoder costs memory and download time, so load only what you need:
```python
model = QuadEmbed.from_pretrained(modalities=("text", "vision")) # skip audio
model = QuadEmbed.from_pretrained(device="cpu") # force CPU
```
Audio takes mono float32 arrays at 16 kHz; video takes a file path
(`model.embed_video_file("clip.mp4")`). A complete runnable example lives in
[the source repo](https://github.com/mithilai/QuadEmbed/blob/main/package/examples/inference_example.py).
This repo hosts the trained weights only; the library itself is published on
[PyPI](https://pypi.org/project/quadembed/) so there is one canonical copy of
the code rather than a duplicate here that can drift out of sync.
## Important caveat: two different vision checkpoints
`checkpoints/vision_projector.pt` (best pure-image retrieval, the round-7
result above) and `checkpoints/video_projector.pt` (best video retrieval)
are **not the same weights**. The video checkpoint was produced by an
earlier training path that continued a *pre-round-3* vision projector
(before this project's spatial-merge architecture fix) on real MSR-VTT
video frames. It was never re-trained against the improved architecture or
the larger round-7 dataset. Practically: use `vision_projector.pt` for
image retrieval, `video_projector.pt` for video retrieval, and don't expect
them to be interchangeable or to represent the same underlying model
version.
## Scope-downs vs. the GELATO paper
- Training data: ~172k image-caption pairs (Flickr8k + COCO + Conceptual
Captions), ~600 audio-caption pairs (AudioCaps), ~180 video-frame pairs
(MSR-VTT) β€” vs. the paper's enterprise-scale, multi-domain corpus.
- Batch size up to 128 vs. the paper's 256.
- No task-specific LoRA adapters (retrieval/classification/clustering)
layered on top β€” base cross-modal alignment only.
- Six rounds of experimentation ruled out architecture (merge fidelity,
projector capacity) as the remaining gap to the paper's published numbers;
data volume was the one lever that reliably helped, and this repo's
training data is still 2-3 orders of magnitude smaller than an
enterprise-scale corpus.
## License
This repo is licensed **CC-BY-NC-4.0** (non-commercial), matching the most
restrictive license among its components:
[`jina-embeddings-v5-text-nano`](https://huggingface.co/jinaai/jina-embeddings-v5-text-nano)
is CC-BY-NC-4.0. `siglip2-base-patch16-naflex` and `whisper-large-v3` are
both Apache-2.0. If you plan to use this for anything commercial, you'll
need a compatible license for the text encoder specifically β€” check with
Jina AI directly.
## Citation
QuadEmbed reproduces the architecture described in Jina AI's GELATO paper.
All credit for the original architecture and training recipe belongs there:
```
@article{gelato2026,
title={jina-embeddings-v5-omni / GELATO},
note={arXiv:2605.08384}
}
```
QuadEmbed is not affiliated with or endorsed by Jina AI β€” an independent,
from-scratch reproduction of their published architecture, built as a
learning and portfolio project.