Publish TE-86M safetensors release

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.gguf filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,227 @@
+---
+language:
+- en
+license: apache-2.0
+tags:
+- multimodal
+- embedding
+- matryoshka
+- trimodal
+- image-text-audio
+- retrieval
+- cross-modal
+- edge
+- rag
+library_name: safetensors
+pipeline_tag: feature-extraction
+datasets:
+- custom
+---
+
+# TE-86M — Trimodal Embeddings (Depth-2)
+
+**TE-86M** maps image, audio, and text into a shared 1280-dim embedding space, enabling cross-modal retrieval with a single vector index. All three modalities share a unified space with full Matryoshka truncation support down to 128 dims.
+
+Built for edge deployment — the entire model runs on a Raspberry Pi 5.
+
+Successor to [TE-75M](https://huggingface.co/augmem/TE-75M), with depth-2 residual projection heads that break through the cross-modal retrieval ceiling of depth-1 architectures while maintaining text retrieval quality.
+
+> Also available in [GGUF format](https://huggingface.co/augmem/TE-86M-GGUF) for quantized edge deployment.
+
+## Architecture
+
+TE-86M uses lightweight edge encoders with depth-2 residual projection heads that expand through a 1920-dim hidden layer before projecting into a shared 1280-dim embedding space:
+
+```
+Text  --> LEAF-IR (768-d) -----------> DeepProjectionHead-d2 (768 -> 1920 -> 1920 -> 1280)
+Image --> MobileNetV4-Medium (1280-d) --> DeepProjectionHead-d2 (1280 -> 1920 -> 1920 -> 1280)
+Audio --> EfficientAT mn20_as (1920-d) --> DeepProjectionHead-d2 (1920 -> 1920 -> 1920 -> 1280)
+```
+
+All outputs are L2-normalized into the shared 1280-dim space for cross-modal cosine similarity.
+
+| Component | Architecture | Params | Size |
+|---|---|---|---|
+| Text encoder | LEAF-IR (MongoDB/mdbr-leaf-ir) | 22.7M | 87.2 MB |
+| Image encoder | MobileNetV4-Medium (timm) | 8.4M | 32.4 MB |
+| Audio encoder | EfficientAT mn20_as | 17.9M | 68.5 MB |
+| Image projection | DeepProjectionHead-d2 (1280 -> 1920 -> 1920 -> 1280) | 12.3M | 47.0 MB |
+| Audio projection | DeepProjectionHead-d2 (1920 -> 1920 -> 1920 -> 1280) | 13.5M | 51.7 MB |
+| Text projection | DeepProjectionHead-d2 (768 -> 1920 -> 1920 -> 1280) | 11.3M | 43.2 MB |
+| **Total** | | **86.1M** | **329.9 MB** |
+
+### Projection head detail
+
+Each `DeepProjectionHead-d2` is a depth-2 residual MLP with Matryoshka-aware training:
+
+```
+Linear(encoder_dim, 1920) -> GELU -> LayerNorm -> Dropout(0.3)
+  -> Linear(1920, 1920) -> GELU -> LayerNorm -> Dropout(0.3) + residual
+  -> Linear(1920, 1920) -> GELU -> LayerNorm -> Dropout(0.3) + residual
+  -> Linear(1920, 1280)
+```
+
+### Why depth-2?
+
+Ablation experiments showed depth-1 heads hit an I->T retrieval ceiling at ~0.60 R@1 regardless of hyperparameter tuning. Depth-2 heads broke through to 0.618, providing the representational capacity to serve cross-modal AND text retrieval simultaneously. The extra 11M params (75M -> 86M) remain edge-viable.
+
+### Matryoshka dimensions
+
+Embeddings can be truncated to `[1280, 768, 512, 256, 128]` dimensions while preserving retrieval quality — trained with Matryoshka Representation Learning (MRL).
+
+## Benchmarks
+
+All benchmarks run on a single NVIDIA L4 GPU with 5K SALT samples.
+
+### Cross-modal retrieval — SALT (5K trimodal samples)
+
+| Direction | TE-86M (86M) | TE-75M (75M) | ImageBind (1.2B) | EBind (1.78B*) |
+|---|---|---|---|---|
+| Image -> Text R@1 | 0.618 | 0.615 | 0.736 | **0.783** |
+| Text -> Image R@1 | 0.630 | 0.614 | 0.712 | **0.779** |
+| Text -> Audio R@1 | **0.108** | 0.103 | 0.038 | 0.047 |
+| Audio -> Text R@1 | 0.087 | 0.082 | 0.039 | 0.035 |
+| Image -> Audio R@1 | **0.068** | 0.062 | 0.023 | 0.027 |
+| Audio -> Image R@1 | **0.070** | 0.063 | 0.025 | 0.032 |
+
+### Audio retrieval — AudioCaps & Clotho
+
+| Benchmark | Direction | TE-86M | TE-75M | CLAP-Large | ImageBind | EBind |
+|---|---|---|---|---|---|---|
+| AudioCaps | A->T R@1 | 0.229 | 0.210 | **0.420** | 0.116 | 0.225 |
+| AudioCaps | T->A R@1 | 0.156 | 0.148 | **0.280** | 0.080 | 0.219 |
+| Clotho | A->T R@1 | **0.219** | 0.208 | 0.195 | 0.061 | 0.088 |
+| Clotho | T->A R@1 | **0.177** | 0.172 | 0.167 | 0.074 | 0.118 |
+
+### Image-text retrieval — MSCOCO & Flickr30k
+
+| Benchmark | Direction | TE-86M (86M) | TE-75M (75M) | EBind (1.78B*) | ImageBind (1.2B) |
+|---|---|---|---|---|---|
+| Flickr30k | I->T R@1 | 0.494 | 0.478 | **0.951** | 0.918 |
+| Flickr30k | T->I R@1 | 0.332 | 0.303 | **0.853** | 0.766 |
+| MSCOCO 5K | I->T R@1 | 0.343 | 0.320 | **0.743** | 0.658 |
+| MSCOCO 5K | T->I R@1 | 0.225 | 0.208 | **0.559** | 0.490 |
+
+### Zero-shot classification — ESC-50
+
+| Model | Params | Accuracy |
+|---|---|---|
+| TE-86M | 86M | **93.9%** |
+| CLAP-Large | 67.8M | 90.5% |
+| TE-75M | 75M | 93.2% |
+| EBind | 1.78B* | 77.0% |
+| ImageBind | 1.2B | 66.4% |
+
+### Text retrieval — MTEB (NDCG@10)
+
+Text-text retrieval quality in the shared embedding space, measured on MTEB retrieval tasks:
+
+| Task | TE-86M | TE-75M | Raw LEAF-IR | Recovery |
+|---|---|---|---|---|
+| ArguAna | 0.545 | 0.544 | 0.594 | 92% |
+| CQADupstackGaming | 0.515 | 0.506 | 0.607 | 85% |
+| CQADupstackUnix | 0.334 | 0.355 | 0.428 | 78% |
+| FEVERHardNegatives | 0.561 | 0.551 | 0.863 | 65% |
+| HotpotQAHardNegatives | 0.554 | 0.531 | 0.700 | 79% |
+| FiQA2018 | 0.291 | 0.292 | 0.392 | 74% |
+| ClimateFEVER | 0.231 | 0.215 | 0.353 | 65% |
+| SCIDOCS | 0.154 | 0.153 | 0.198 | 78% |
+| TRECCOVID | 0.507 | 0.474 | 0.820 | 62% |
+
+TE-86M improves MTEB text retrieval over TE-75M on 7/9 tasks. The depth-2 projection heads recover 62-92% of raw LEAF-IR's retrieval quality while mapping into the cross-modal shared space.
+
+## Usage
+
+### Loading components
+
+```python
+from safetensors.torch import load_file
+
+# Load entire model
+tensors = load_file("TE-86M.safetensors")
+
+# Extract components by prefix
+text_enc_sd = {k.removeprefix("text_encoder."): v for k, v in tensors.items() if k.startswith("text_encoder.")}
+image_enc_sd = {k.removeprefix("image_encoder."): v for k, v in tensors.items() if k.startswith("image_encoder.")}
+audio_enc_sd = {k.removeprefix("audio_encoder."): v for k, v in tensors.items() if k.startswith("audio_encoder.")}
+image_proj_sd = {k.removeprefix("image_projection."): v for k, v in tensors.items() if k.startswith("image_projection.")}
+audio_proj_sd = {k.removeprefix("audio_projection."): v for k, v in tensors.items() if k.startswith("audio_projection.")}
+text_proj_sd = {k.removeprefix("text_projection."): v for k, v in tensors.items() if k.startswith("text_projection.")}
+```
+
+### Matryoshka truncation
+
+```python
+import torch.nn.functional as F
+
+# Full 1280-dim embedding
+embedding = model(input)  # (N, 1280)
+
+# Truncate to 256-dim and re-normalize
+embedding_256 = F.normalize(embedding[:, :256], dim=-1)
+```
+
+## File layout
+
+```
+TE-86M.safetensors     # All components in one file (~330 MB)
+```
+
+### Tensor key prefixes
+
+| Prefix | Component | Tensors |
+|---|---|---|
+| `text_encoder.*` | LEAF-IR (float32) | 103 |
+| `image_encoder.*` | MobileNetV4-Medium | 462 |
+| `audio_encoder.*` | EfficientAT mn20_as | 312 |
+| `image_projection.*` | Depth-2 projection head | 14 |
+| `audio_projection.*` | Depth-2 projection head | 14 |
+| `text_projection.*` | Depth-2 projection head | 14 |
+
+## Training
+
+- **Loss**: InfoNCE (contrastive) with Matryoshka Representation Learning
+- **Data**: ~2.2M synthetically generated trimodal triplets (WordNet) + 200K MSCOCO img+txt + 262K WavCaps aud+txt + 1.5M Nomic text pairs
+- **Hardware**: 2x NVIDIA L4 GPUs
+- **Optimizer**: AdamW, lr=1.41e-3, weight decay=1e-4, cosine scheduler
+- **Epochs**: 50
+- **Batch size**: 4096
+- **Dropout**: 0.20 -> 0.25 (ep27) -> 0.30 (ep29) — mid-run regularization increases
+- **Text mixing**: λ_tt=0.5 (ep1-9) -> 0.25 (ep10-50) — Nomic supervised text pairs
+- **Projection heads only** — source encoders are frozen during training
+
+### Improvements over TE-75M
+
+| Change | TE-75M | TE-86M |
+|---|---|---|
+| Projection depth | 1 (single residual block) | 2 (two residual blocks) |
+| Head params | 26.1M | 37.2M |
+| Total params | 75.2M | 86.1M |
+| SALT I->T R@1 | 0.615 | 0.618 (+0.5%) |
+| SALT T->I R@1 | 0.614 | 0.630 (+2.6%) |
+| MSCOCO I->T R@1 | 0.320 | 0.343 (+7.2%) |
+| Clotho A->T R@1 | 0.208 | 0.219 (+5.3%) |
+| ESC-50 | 93.2% | 93.9% (+0.7%) |
+
+### Design decisions
+
+- **Depth-2 residual heads**: Ablation confirmed depth-1 hits I->T ceiling at ~0.60 regardless of dropout or λ_tt. Depth-2 provides capacity to serve cross-modal and text retrieval simultaneously.
+- **3-head shared space**: All modalities project into a learned 1280-dim space (image-native dimension)
+- **LEAF-IR text encoder**: 23M-param retrieval-optimized text encoder enables fully edge-deployable text inference
+- **Frozen source encoders**: MobileNetV4, EfficientAT, and LEAF-IR are kept frozen; only projection heads are trained
+- **Edge-first**: All source encoders can run on devices like Raspberry Pi 5
+
+## Limitations
+
+- Audio retrieval lags behind specialist models like CLAP on audio-only benchmarks
+- Image-text retrieval trades accuracy vs larger vision encoders for edge deployability
+- Text retrieval recovers 62-92% of raw LEAF-IR quality (gap is domain-dependent)
+
+## Links
+
+- **Website**: [augmem.ai](https://augmem.ai)
+- **GitHub**: [github.com/augmem](https://github.com/augmem)
+
+## License
+
+Apache 2.0

TE-86M.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5245f12b721086d90b4fe649c8f38b6928e658ba81087a620398f3dec567e2b7
+size 346056964