Skip to content

Latest commit

 

History

15 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

QuadEmbed

PyPI Python License: CC BY-NC 4.0 Model on HF

A local, from-scratch-trained multimodal embedding model covering all four modalities, text, image, audio, and video, in one shared 768-dim embedding space. Trained end-to-end on a single RTX 4060 laptop GPU (8GB VRAM).

Reproduces the architecture behind jina-embeddings-v5-omni ("GELATO": Geometry-preserving Embeddings via Locked Aligned TOwers, arXiv 2605.08384), built on public source encoders instead of Jina's exact adapted checkpoints. Full credit for the architecture and training recipe goes to Jina AI's GELATO paper. QuadEmbed is an independent, from-scratch reproduction of it, not a copy of Jina's released weights, and is not affiliated with or endorsed by Jina AI.

Model weights huggingface.co/Mithil-AI/quadembed-nano
Python package pypi.org/project/quadembedpip install quadembed
Write-up I Built a Multimodal Embedding Model From Scratch on an RTX 4060
Experiment log EXPERIMENTS.md — seven rounds, what moved the needle and what didn't

Install

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()   # pulls weights from the Hub
text = model.embed_text(["a dog running on the beach"])
image = model.embed_image([Image.open("photo.jpg").convert("RGB")])
similarity = text @ image.T           # L2-normalized, so this is cosine

The rest of this README covers reproducing the training from source. The packaged library lives in package/.

Results

Cross-modal retrieval recall@k on held-out splits, text-query direction:

Modality R@1 R@5 R@10 n Peak VRAM
Image 13.7% 68.6% 81.1% 1024 1.56 GB
Audio 67% 97% 100% 33 2.69 GB
Video (held-out) 40% 86% 94% 50 0.89 GB

Random-chance R@1 on the 1024-candidate image eval is ~0.1%. Every training run stayed under 2.7GB peak VRAM, well inside the 8GB budget.

Architecture

Three frozen encoders, two small trainable projectors:

Role Model (frozen) Output dim
text jinaai/jina-embeddings-v5-text-nano (239M) 768 (joint space, the anchor)
vision google/siglip2-base-patch16-naflex 768/patch, merged to 3072
audio openai/whisper-large-v3 encoder only 1280/frame

Trainable: VisionProjector.fc_vision_2 (3072→768, 2.36M params) and AudioProjector.fc_audio (1280→768, 0.98M params), plus small per-modality delimiter vectors. Everything else stays frozen and untouched, matching GELATO's own ablation showing this outperforms unfreezing.

Video gets no separate encoder: quadembed/video.py samples frames from an mp4 and runs each through the same trained vision projector, then mean-pools over time.

Why these substitutions: the paper's vision/audio encoders (Qwen3.5's vision tower, Qwen2.5-Omni's audio tower) aren't distributed as standalone checkpoints and are described in the paper as adapted from SigLIP2 and Whisper-large-v3 respectively, so this reproduction uses those source models directly. The text encoder is an exact match.

Datasets

All public, all pulled from the Hugging Face Hub at training time. No proprietary or manually collected data is used anywhere in this project.

Dataset Modality Used for Size used here License
jxie/flickr8k image-text vision training + held-out eval ~30,000 pairs (6k images × 5 captions) CC BY 4.0
jxie/coco_captions image-text vision training ~124,560 pairs (40 of 182 shards) CC BY 4.0 (COCO)
google-research-datasets/conceptual_captions image-text vision training (domain diversity) 17,291 images Google CC BY (see note)
OpenSound/AudioCaps audio-text audio training + eval ~600 clip-caption pairs MIT (AudioCaps: CC BY 4.0)
VLM2Vec/MSR-VTT video-text video training + held-out eval 60 train / 50 test clips research use

Total vision training set at the final round: ~172,000 image-caption pairs across three datasets and two visual domains (curated photography from Flickr8k/COCO, plus web imagery from Conceptual Captions).

Notes on specific datasets:

  • Conceptual Captions (CC3M) ships only (image_url, caption) pairs, not image bytes. scripts/fetch_cc3m.py downloads the actual images concurrently to local disk (32 workers, 6s timeout, skip-on-failure) so training never depends on live network calls. Expect ~58% success; dead links and expired stock-photo hosts account for the rest. It also filters degenerate images (min(w,h) < 32 or aspect ratio beyond 6:1) — two one-pixel-wide "images" produced NaN in SigLIP2's patch-grid math and silently poisoned an entire training run before this guard existed.
  • AudioCaps is fetched by parquet shard path rather than streaming; streaming the same data measured ~200x slower due to per-row-group HTTP range requests against these audio-heavy files. Audio is decoded manually via soundfile rather than the default torchcodec path, which fails to load on Windows/CUDA (WinError 127).
  • COCO Captions is fetched in chunks of 4 shards and concatenated; a single bulk fetch stalled around 6 shards.
  • Held-out evaluation always uses the Flickr8k test split for vision, regardless of what was added to training, so numbers stay comparable across all seven experiment rounds.

Repository layout

src/
  quadembed/           importable library
    encoders.py        frozen text/vision/audio encoder wrappers
    projectors.py      the two trainable projectors + video pooling
    losses.py          InfoNCE + Matryoshka contrastive loss
    data.py            dataset loaders for all five datasets above
    video.py           frame sampling + video embedding
  scripts/             command-line entrypoints
    train.py           train one modality's projector
    eval.py            recall@k for vision/audio
    eval_video.py      recall@k for video
    fetch_cc3m.py      one-time CC3M image downloader
    smoke_test.py      verify encoders load and shapes line up
hf_repo/               staging folder for the Hugging Face release
EXPERIMENTS.md         full seven-round experiment log

Setup

python -m venv .venv
.venv\Scripts\Activate.ps1
pip install -r requirements.txt
cd src

Verify the frozen encoders load and tensor shapes line up before training:

python -m scripts.smoke_test

Training

Sequential, one modality at a time (per GELATO's constrained-hardware guidance: only one frozen encoder needs to be resident at a time). Run from the src/ directory:

# vision, final configuration (~172k pairs)
python -m scripts.train --modality vision --steps 9000 --batch_size 128 --num_workers 4 --coco_shards 40 --cc3m_dir ../data_cache/cc3m

# audio
python -m scripts.train --modality audio --steps 500 --batch_size 16 --num_workers 2 --num_shards 6

# video (continues the vision projector on real sampled video frames)
python -m scripts.train --modality video --steps 300 --num_videos 60 --frames_per_video 3

To use Conceptual Captions, download the images first:

python -m scripts.fetch_cc3m --num_urls 30000 --workers 32 --out_dir ../data_cache/cc3m

--num_workers 4 matters for vision: SigLIP2's NaFlex preprocessing is CPU-bound enough that num_workers=0 drops GPU utilization to near zero. Continue any run from an existing checkpoint with --init_checkpoint.

Evaluation

python -m scripts.eval --modality vision --checkpoint ../checkpoints/vision_projector_FINAL_round7_best.pt
python -m scripts.eval --modality audio --checkpoint ../checkpoints/audio_projector_step500.pt
python -m scripts.eval_video --checkpoint ../checkpoints/video_projector_step300.pt --num_videos 50

Reports text→modality and modality→text recall@{1,5,10} on a held-out split.

What seven rounds of experiments actually showed

The short version: architecture was never the bottleneck, data volume was. Rounds 3 through 6 fixed a genuine spatial-merge bug, doubled batch size, added a genuinely different visual domain, and increased projector capacity. None of them moved R@1 off ~13%. Round 7 scaled the training set 3.2x and moved every metric at once.

Full round-by-round detail, including the negative results and the NaN incident, is in EXPERIMENTS.md.

Known scope-downs vs. the paper

  • Batch size 16-128 vs. the paper's 256; in-batch negative count matters a lot for InfoNCE quality.
  • ~172k image-caption pairs vs. the paper's enterprise-scale, multi-domain corpus. Confirmed empirically to be the actual bottleneck.
  • No task-specific LoRA adapters (retrieval/classification/clustering) layered on top; base cross-modal alignment only.

License

Code in this repo is available for research and educational use. Note that the frozen text encoder (jina-embeddings-v5-text-nano) is CC-BY-NC-4.0 (non-commercial), which carries through to any model trained on top of it. SigLIP2 and Whisper-large-v3 are Apache-2.0. Check each dataset's own license before redistribution.

About

No description, website, or topics provided.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages