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

from mind_standalone import load_mind, embed   # github.com/taylor-geospatial/mind
emb = 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/mind

MIND-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.

from mind_standalone import load_mind, embed
m = load_mind("mind_small.safetensors")
emb = 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.