Distilling a 33M model to 9.8M, then making the small model brutally fast with a pure-Rust inference engine (AMX + BNNS + NEON). No GPU, no Python, no warm-up. 5.4 ms cold start.
Embedding models face a quiet but brutal constraint: the smaller they get, the less they know — and the bigger they get, the harder they are to deploy on the edge. Most teams resolve this by renting a GPU or shipping the model inside a Python runtime that needs 29 seconds to even start up.
We went the other way. We distilled a 33M teacher into a 9.8M trunk (Hybrid Distillation), then rebuilt the inference stack in pure Rust so that on an Apple M4 CPU the query encoder runs at 103 µs — 53.9× faster than its teacher — with a 5.4 ms cold start, and zero Python at runtime.
This post is a technical walkthrough of both halves: how we compressed the model, and how we then squeezed every last µs out of the hardware. If you've ever fought with slow cold starts, fp16 promises that didn't materialize, or int8 quant that made things slower — the second half is for you.
TL;DR
| VelaVec (ours) | Teacher (bge-small) | Ratio | |
|---|---|---|---|
| Parameters | 9.8M | 33M | 30% |
| Embedding dim | 256 | 384 | — |
| Query latency (M4 CPU) | 103 µs | 5.5 ms | 53.9× |
| Doc latency (128 tok, CPU) | ~580 µs (2-thread) | 13 ms | ~22× |
| Cold start | 5.4 ms (Rust) | ~29 s (PyTorch) | ~2600× |
| Weight size | 39 MB fp32 | ~130 MB | 30% |
Single trunk, dual pooling heads: head=0 for symmetric semantics (STS, clustering, dedup), head=1 for asymmetric retrieval (query ↔ doc). Two use-cases from one 39 MB file.
Part 1 — Compressing to 9.8M without losing the teacher's knowledge
A naive student would just copy the teacher's outputs and fail to generalize. We used a hybrid architecture that treats the vocabulary table as a frozen memory bank and buys capacity with a shallow attention stack instead of parameter count.
Architecture: frozen table + 3 attention layers
tokens ──▶ frozen 30,522×256 static table (distilled) ──▶ 3× BiBlock
(RMSNorm, RoPE, SwiGLU)
──▶ attention pooling ──▶ 256-d
Three choices matter:
-
Frozen distillation table. Instead of learning a fresh token embedding matrix, we distill the teacher's token embeddings into a 30,522×256 static table (PCA-compressed from bge-small). This is a dense, pre-packed source of semantic prior — the model starts every sentence knowing roughly what every word means, before any attention happens.
- Ablation: zeroing the Engram/N-gram memory component costs +32.7% val_loss. The table is not decoration; it carries load.
- Ablation: training the same-size student from scratch with 7× more parameters still lands 6.4 recall points lower. The frozen table is where the compression actually lives.
Three layers of bidirectional attention (RMSNorm + RoPE + SwiGLU FFN,
ffn_mult=2). We swept 2 vs 3 layers: with hard negatives + the full dataset + attention pooling, 3 layers wins at the same parameter budget. The FFN/q-proj shapes are chosen so the whole model fits a 256-d embedding contract.Dual heads, one trunk. The trunk learns representations; two additive pooling heads specialize them per task family. This is how one 9.8M model carries both symmetric and retrieval semantics without doubling the file size.
Training protocol: 5-domain distillation
Distilled from bge-small (33M) across five domains with a multi-task objective:
- Relational KD — the student reproduces the teacher's pairwise relation (similarity structure) on sampled pairs.
- InfoNCE / retrieval contrastive loss — hard negatives from MARCO-style retrieval pools.
- Embedding regression — the student matches the teacher's absolute embedding where it matters (retrieval head).
- Tasks: NLI (sym), STS12–17 (sym), MS MARCO (retrieval), SciFact/NFCorpus/ArguAna (retrieval).
Loss shaping details turned out to matter as much as the recipe: cosine annealing with warm-up, dropout, and per-250-step early stopping were required to avoid the classic "peak at step ~500 then collapse" failure mode.
Where the 9.8M lands on real eval
MTEB-style (no prompts, max_len 256, CPU):
| Task | VelaVec | bge-small | % of teacher |
|---|---|---|---|
| STS12 (cosine spearman) | 0.7117 | 0.7744 | 91.9% |
| STSBenchmark | 0.7596 | 0.8586 | 88.4% |
| Banking77 (kNN acc) | 0.8365 | 0.8175 | 102.3% |
| SciFact (ndcg@10) | 0.6310 | 0.7200 | 87.6% |
| ArguAna (ndcg@10) | 0.4560 | 0.5950 | 76.6% |
| NFCorpus (ndcg@10) | 0.2449 | 0.3371 | 72.7% |
Internal pools: NLI r@1 = 0.912, MS MARCO dev r@1 = 0.694 (1k-candidate pool).
The honest read: we're at parity-or-better where the domain matches (Banking77 beats the teacher; symmetric tasks ~92%), and ~73–88% on OOD retrieval. That's the price of 30% of the parameters — but the deployment story below shows why we think it's a good trade.
Part 2 — Making a small model run fast: the pure-Rust inference engine
This is the part that took the longest and taught us the most. The model is small; why isn't it instant? The answer was: the framework was in the way. Our Rust engine (BNNS/CoreML pre-packed GEMM graphs + NEON kernels) eventually hit ~2.4–3.0 TF/s single-thread AMX peak on Apple silicon — vs ~35–40% of peak for BLAS. Here's the journey, including the dead ends.
1. Measure the machine first, then the model
Anything we measured was useless until we controlled the machine state. Single-item latency on the same code swings 2× depending on system load:
- 414 µs @ load ~17 → 218–226 µs @ load ~2.5 → 165 µs @ load < 1
Rule we now follow: same-session A/B only. Cross-session numbers quietly lie.
2. The fp16 trap on Apple CPU
We assumed half-precision would be free speed — it isn't. Apple AMX gives no fp16 acceleration for small GEMMs, and fp16 was slower than fp32 in our case (0.74–0.90×). Verdict: fp32 binds to AMX, fp16 doesn't.
3. int8 didn't help either
BNNSMatMul/filter reject Int8 outright. Via CoreML+BNNSGraph:
- int8 weights ≈ fp32 speed (no win)
- int8 weights + activations ≈ 2× slower
There is no public AMX int8 GEMM path (verified at both l=15 and l=256, compute-bound). Quantization is not free speed on Apple silicon — a lesson that saved us from shipping a regression.
4. The big win: BNNSGraph fp32 with weight pre-packing
The single biggest speedup came from an unglamorous thing: letting CoreML pre-pack the weights and using BNNSGraph instead of cblas_sgemm:
- At the model's actual GEMM shapes (inner dim l=15): 159 µs → 51 µs for 15 GEMMs — 2.5–3.1× over BLAS
- AMX peak: ~2.4–3.0 TF/s single-thread vs 35–40% for BLAS
Pre-packing means the weight layout is optimized once at load time; every forward pass after that pays nothing for it.
5. NEON elementwise kernels: attention to scalar paths
The paper-tiger GEMMs were solved; the elementwise ops were next. We replaced vvexpf/scalar paths with hand-written NEON kernels:
- rmsnorm: dual-chain reduction
-
SwiGLU + feed-forward gate: fused, with a degree-9
exp2polynomial (coeffs from numpy polyfit on2^(u/2), u∈[-1,1], max rel error 5.5e-14) instead of a generic exp — this matters: a degree-7 poly (4.7e-7 err) flipped 1 in 500 recall queries; degree-9 restored r@1 exactly. - RoPE: 4-wide vectorized.
4-wide is the sweet spot on Apple's OoO cores; 8-wide interleaved versions add register pressure and don't go faster (the hardware already extracts ILP between iterations).
6. Cold start: 5.4 ms vs 29 s
PyTorch's 29 s cold start on the M4 is mostly framework import + weight deserialization. Our Rust engine: 5.4 ms — the model fits comfortably in L2-ish footprint, weights are memory-mapped once, and there's no interpreter to boot. For serverless / autoscale-any-request workloads this is the difference between "add an instance" and "just add the model".
Part 3 — Multimodal extension: VelaVec-T2I
The text trunk is fast enough that adding a tiny vision bridge costs almost nothing. VelaVec-T2I appends a projection head (not a vision tower): frozen CLIP ViT-B/32 (512d) + ViT-L/14 (768d) images are concatenated (1280d) and mapped by a residual MLP (1280→768→256) into VelaVec space — ±1.5M extra parameters.
- text→img r@1 (flickr30k): 0.560 (teacher ensemble: 0.585)
- img→text r@1: 0.619
- The vision head is projection-only — no CLIP weights are bundled; teacher features come from HF transformers at inference.
One nuance worth sharing: the best img→text head and the best text→img head are different projection heads. Regression terms pull the image embedding toward caption-space centroids, which helps i2t recall but hurts the cross-image separability that t2i ranking needs. We shipped both variants — a 2 GB class of lesson in "what you optimize for is what you get."
Repo / models
- VelaVec (pure text, Rust inference): https://huggingface.co/xagent2025/VelaVec
- VelaVec-T2I (text→image retrieval head): https://huggingface.co/xagent2025/VelaVec-T2I
- All weights are self-contained
model.safetensors(fp32); Rustweights.binfor the edge path.
When is this the right tool?
Use VelaVec when your bottleneck is deployment, not accuracy ceiling:
- Edge / mobile — 39 MB, 5.4 ms cold start, pure-Rust, no Python runtime
- Serverless & bursty — every request can cold-start its own instance almost free
- High-QPS catalog / dedup — 103 µs × millions of items is a different cost model than GPU batch
- Embedded DBs (e.g. MongoDB $vectorSearch-style) — in-process encoding, no RPC hop, zero-cost incremental writes
Don't use it when you need full bge-large/GTE-class OOD robustness and have the GPU budget — that's a different product.
Model: hybrid distillation + dual heads. Inference: Rust + BNNS/CoreML pre-packed graphs + NEON. Benchmarked on Apple M4 CPU, fp32, single thread unless noted. Numbers are same-session A/B under controlled machine load.
Top comments (0)