README.md

---
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
```python
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.

```python
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.