Hugging Face's logo

isaaccorley
/
MIND

ONNX
location-encoder
geospatial
earth-observation
Model card Files
xet
 Community
MIND / README.md
isaaccorley's picture
isaaccorley
feat: add MIND-small (distilled 6.4M student) weights + deploy graphs + card section
43bd70f verified
preview code
|
Raw
History Blame Contribute Delete
4.7 kB
metadata 
license: mit
tags:
  - location-encoder
  - geospatial
  - earth-observation

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.