MIND
Matryoshka Implicit Neural Distillation: a lat/lon location encoder distilled from four geospatial
foundation models (AlphaEarth, Climplicit, GeoCLIP, SINR). Given a coordinate it returns an embedding
with no imagery or labels at inference; the embedding is Matryoshka-structured, so a 64-d prefix is the
deploy width.
Files
mind.safetensors - recommended weights, fp16 (227 MB, safe / no-pickle). Load with mind_standalone.load_mind.mind.pt - fp32 weights (454 MB) for exact reproduction.mind.onnx (+ mind.onnx.data) - ONNX graph for onnxruntime / TensorRT (input latlon [N,2], dynamic batch).mind.pt2 - torch ExportedProgram; load with torch.export.load.mind_small.safetensors / mind_small.onnx / mind_small.pt2 - MIND-small, the distilled 6.4 M student (see below).
Usage
1from mind_standalone import load_mind, embed # github.com/taylor-geospatial/mind
2emb = embed(load_mind("mind.safetensors"), lats, lons) # [N, 64]; half=True for ~9x on GPU
Build a TensorRT engine from the ONNX with
trtexec --onnx=mind.onnx --saveEngine=mind.trt --fp16.
Code and paper:
https://github.com/taylor-geospatial/mindMIND-small (distilled, deployable)
A small student self-distilled from MIND's 128-d prefix: a width-1024, depth-6 ReSIREN (6.4 M params,
18x smaller, ~18x fewer FLOPs) trained to match the teacher's first 128 trunk dims (the frozen
teacher generates targets on sampled coordinates, so there is no dataset). It emits a 128-d embedding.
1from mind_standalone import load_mind, embed
2m = load_mind("mind_small.safetensors")
3emb = embed(m, lats, lons, dim=128, feature="head") # [N, 128]; the head output IS the embedding
- Fidelity: 0.99 cosine to the teacher's 128-d prefix.
- Downstream: retains ~92% of the teacher prefix's macro transfer on the spatial-CV linear-probe
suite -- and beats the full 3072-d trunk, which overfits under spatial holdout.
- Throughput: ~11x the teacher on GPU (TensorRT fp16: 29.5M pts/s on an H100) and ~14x on CPU
(79k vs 5.6k pts/s on 16 cores), at ~half the memory. Makes CPU-only / edge inference practical.
| teacher (3072) | MIND-small (128) |
|---|
| params | 113.5 M | 6.4 M |
| H100 TensorRT fp16 | 2.64M pts/s | 29.5M pts/s |
| 16-core CPU (torch) | 5.6k pts/s | 79k pts/s |
| cosine to teacher[:128] | 1.00 | 0.99 |
Deployment: throughput & memory (which format to use)
500k coordinates, batch 65,536, full 3072-d trunk output. Inputs/outputs are GPU-resident (ONNX/TensorRT
use IOBinding), so throughput reflects compute, not host<->device transfer. Peak memory is the NVML
device maximum on the H100.
| Format | Precision | A100 | H100 | Peak mem (H100) |
|---|
| ONNX -> TensorRT | fp16 | 0.84M | 2.64M pts/s | 5.5 GB |
| torch (autocast) | fp16 | 0.70M | 1.99M pts/s | 3.8 GB |
| torch (torchao) | fp8 e4m3 | n/a | 0.93M pts/s | 6.2 GB |
| ONNX Runtime (CUDA EP) | fp32 (TF32) | 0.39M | 0.91M pts/s | 10.6 GB |
| torch (eager) | fp32 | 0.08M | 0.21M pts/s | 5.0 GB |
ExportedProgram (.pt2) | fp32 | 0.08M | 0.21M pts/s | 4.2 GB |
What to use
- Max throughput / production serving -> ONNX + TensorRT fp16. Fastest on both GPUs; pay a one-time
engine build.
- PyTorch app -> torch fp16 (
half=True). Lowest memory, zero extra deps, ~25% behind TensorRT.
- Outside Python / no TensorRT build -> ONNX Runtime (CUDA EP). Portable and cross-language.
- Torch graph without the package -> ExportedProgram (
torch.export.load("mind.pt2")).
- Tight compute / CPU / edge -> MIND-small (above): ~half the memory, an order of magnitude faster.
Notes on precision:
- fp16 is the sweet spot: ~1% change to the 64-d embedding, large speedup, no extra work.
- fp8 is not worth it here. Off-the-shelf dynamic-activation fp8 (torchao, H100) ran slower than
fp16 -- the model is small (113M params), so per-layer quant overhead isn't amortized -- and dropped
trunk cosine to ~0.96 because the SIREN
sin activations are quantization-sensitive. It would need
quant-aware distillation to be competitive, and TensorRT fp16 already wins without it.
- Low-rank factorization does not help: the 12 wide SIREN blocks are near-full-rank (they encode
high-frequency content), so any FLOP-saving truncation wrecks the embedding. Distillation (MIND-small)
is the compression that works.
- The CUDA-EP "fp32" row uses TF32 tensor-core GEMMs (ORT default), hence ~5x over torch eager fp32.
Reproduce with scripts/bench_deploy.py (A100/H100) in the repo.
