olmoearth-paper-embeddings

# OlmoEarth — Foundation-Model Embeddings for Paper Table 2 This dataset contains pre-extracted embeddings from **26 Earth-observation foundation models** evaluated on the **24 downstream tasks** that make up Table 2 of the OlmoEarth paper: > *OlmoEarth: Stable Latent Image Modeling for Multimodal Earth Observation* > AI2, 2025. arXiv:2511.13655. For every supported `(model, task)` pair we ran the model's encoder over the task's train / validation / test splits with the **paper-best hyperparameters** (val-best from the original sweep, captured in `eval_settings/`). Each split is saved to disk as a single `.pt` file so downstream KNN / linear-probe / fine-tuning experiments can skip re-extraction. Total size on disk is ~1.1 TB. ## Layout ``` {model_group}/ {task}/ train.pt valid.pt test.pt ``` Each `.pt` file is a Python dict produced by `torch.save`: ```python {"embeddings": torch.Tensor[bfloat16], "labels": torch.Tensor[int64]} ``` - **Classification tasks**: `embeddings` is `(N, D)`. `labels` is `(N,)` for single-label tasks, `(N, num_classes)` for the multilabel `m_bigearthnet`. - **Segmentation tasks**: `embeddings` is `(N, H, W, D)` at the model's patch grid (e.g. `(N, 16, 16, 768)` for a 64×64 input with stride 4). `labels` is at *full* image resolution, e.g. `(N, 64, 64)` or `(N, 256, 256)` — the LP head upsamples patch logits to label resolution. - `D` varies by model (128 for the smallest models, up to 4096 for DINOv3-7B). ## Loading ```python import torch d = torch.load("olmoearth_base/m_eurosat/train.pt") embeddings = d["embeddings"].float() # cast bf16 → fp32 labels = d["labels"] print(embeddings.shape, labels.shape) # torch.Size([2000, 768]) torch.Size([2000]) ``` A more complete example that runs KNN and matches paper Table 2 numbers is in `load_example.py`. ## Row ordering Within any single `.pt` file, row *i* of `embeddings` corresponds to row *i* of `labels` — the pair is always consistent. All three splits are saved in **canonical dataset order** (the dataset's native sample enumeration, identical across models). Concretely: - For a given task, `train.pt` row *i* is the **same source sample** in every model's dump, so you can compare / concatenate / ensemble embeddings across models by row index. The same holds for `valid.pt` and `test.pt`. - The order also matches the underlying eval dataset, so you can join row *i* back to that dataset's *i*-th sample. > Note: an earlier revision of this dataset shuffled `train` (only) with > an unseeded RNG, so train rows did **not** line up across models or with > the source dataset (val/test were always ordered). If you downloaded > before this note was added, re-pull the `train.pt` files. KNN/LP results > are unaffected either way — both are row-order invariant. The one exception is AnySat segmentation — see caveat 3 below. ## Models (26 groups) | `model_group` directory | Paper name (Table 2) | Reference | |---|---|---| | `anysat` | AnySat ViT Base | Astruc et al., 2024 | | `clay_large` | Clay ViT Large | Clay v1, 2024 | | `copernicusfm` | CopernicusFM ViT Base | 2025 | | `croma_base` | CROMA ViT Base | Fuller et al., 2023 | | `croma_large` | CROMA ViT Large | Fuller et al., 2023 | | `dino_v3_dinov3_vitb16` | DINOv3 ViT Base | Meta, 2025 | | `dino_v3_dinov3_vitl16` | DINOv3 ViT Large | Meta, 2025 | | `dino_v3_dinov3_vith16plus` | DINOv3 ViT Huge+ | Meta, 2025 | | `dino_v3_dinov3_vit7b16` | DINOv3 ViT 7B | Meta, 2025 | | `dino_v3_dinov3_vitl16_sat` | DINOv3 Sat ViT Large | Meta, 2025 | | `dino_v3_dinov3_vit7b16_sat` | DINOv3 Sat ViT 7B | Meta, 2025 | | `galileo_nano` | Galileo ViT Nano | Tseng et al., 2025 | | `galileo_tiny` | Galileo ViT Tiny | Tseng et al., 2025 | | `galileo_base` | Galileo ViT Base | Tseng et al., 2025 | | `panopticon` | Panopticon ViT Base | 2024 | | `presto` | Presto | Tseng et al., 2023 | | `prithvi_v2_Prithvi-EO-2.0-300M` | Prithvi v2 ViT Large (300M) | NASA-IBM, 2024 | | `prithvi_v2_Prithvi-EO-2.0-600M` | Prithvi v2 ViT Huge (600M) | NASA-IBM, 2024 | | `satlas_base` | Satlas Swin Base | AI2, 2023 | | `terramind_base` | TerraMind ViT Base | IBM, 2025 | | `terramind_large` | TerraMind ViT Large | IBM, 2025 | | `tessera` | TESSERA | 2025 | | `olmoearth_nano` | OlmoEarth ViT Nano | This paper | | `olmoearth_tiny` | OlmoEarth ViT Tiny | This paper | | `olmoearth_base` | OlmoEarth ViT Base | This paper | | `olmoearth_large` | OlmoEarth ViT Large | This paper | Not every model produced embeddings for every task — some models are modality-restricted (e.g. DINOv3 is S2-only, Tessera is only the multi-modal time-series tasks). Missing combinations are simply absent from disk. ## Tasks (24) Order matches columns 1..24 of paper Table 2. | Task directory | Paper name | Modality | Eval method | Metric | |---|---|---|---|---| | `m_bigearthnet` | m-bigearthnet (geobench) | S2 | KNN | µF1 (multilabel) | | `m_so2sat` | m-so2sat (geobench) | S2 | KNN | accuracy | | `m_brick_kiln` | m-brick-kiln (geobench) | S2 | KNN | accuracy | | `m_forestnet` | m-forestnet (geobench) | L8 | KNN | accuracy | | `m_eurosat` | m-eurosat (geobench) | S2 | KNN | accuracy | | `breizhcrops` | BreizhCrops | S2 (time series) | LP | accuracy | | `cropharvest_Peoples_Republic_of_China_6_sentinel1` | CropHarvest-PRC | S1 (time series) | LP | accuracy | | `cropharvest_Peoples_Republic_of_China_6` | CropHarvest-PRC | S2 (time series) | LP | accuracy | | `cropharvest_Peoples_Republic_of_China_6_sentinel1_sentinel2` | CropHarvest-PRC | S1+S2 (time series) | LP | accuracy | | `cropharvest_Togo_12_sentinel1` | CropHarvest-Togo | S1 (time series) | LP | accuracy | | `cropharvest_Togo_12_sentinel2` | CropHarvest-Togo | S2 (time series) | LP | accuracy | | `cropharvest_Togo_12_sentinel2_sentinel1` | CropHarvest-Togo | S1+S2 (time series) | LP | accuracy | | `m_cashew_plant` | m-cashewplant (geobench) | S2 | LP | mIoU | | `m_sa_crop_type` | m-SA-crop-type (geobench) | S2 | LP | mIoU | | `pastis_sentinel1` | PASTIS | S1 (time series) | LP | mIoU | | `pastis_sentinel2` | PASTIS | S2 (time series) | LP | mIoU | | `mados` | MADOS | S2 | LP | mIoU | | `sen1floods11` | Sen1Floods11 | S1 | LP | mIoU | | `awf_landsat` | AWF (real-world) | L8 (time series) | KNN | accuracy | | `awf_sentinel1` | AWF (real-world) | S1 (time series) | KNN | accuracy | | `awf_sentinel2` | AWF (real-world) | S2 (time series) | KNN | accuracy | | `nandi_landsat` | Nandi (real-world) | L8 (time series) | KNN | accuracy | | `nandi_sentinel1` | Nandi (real-world) | S1 (time series) | KNN | accuracy | | `nandi_sentinel2` | Nandi (real-world) | S2 (time series) | KNN | accuracy | A few additional tasks (`pastis_sentinel1_sentinel2`, the `pastis128_*` variants) are present on disk for some models but were *not* reported in paper Table 2. ## Reproducing the paper The directory `eval_settings/` contains the per-`(model, task)` paper-best hyperparameters used during extraction: - `max_eval_settings_per_group_merged.enriched.json` — external FMs - `{nano,tiny,base,large}_settings.enriched.json` — OlmoEarth checkpoints Each entry records the val-best `pooling_type`, `probe_lr`, the original `run_id`, and an `enriched` `norm_mode` (`pre_trained` vs `dataset`) derived by reading the paper W&B run's actual `--model.use_pretrained_normalizer` flag. We needed `norm_mode` because the JSON's `norm_stats_from_pretrained` field is uninformative for several models (Galileo, Tessera, …) whose helper hardcodes that flag regardless of which sweep arm won. To reproduce paper-Table-2 numbers from these embeddings: - **KNN tasks** (single-timestep classification): paper-Table-2 KNN runs are deterministic given the embeddings — see `load_example.py`. - **LP tasks**: train the linear probe with `probe_lr` from the JSON, `probe_batch_size` from `EVAL_TASKS` in `olmoearth_pretrain.internal.all_evals`, **`linear_probe_eval_interval=5`** (paper used `--select_best_val`), `epochs=50`, and `select_final_test_miou_based_on_epoch_of_max_val_miou=True`. We've verified end-to-end that the saved embeddings reproduce paper Table 2 to within ±0.5–1 pt for >90% of cells, with documented exceptions listed below. ## Known caveats A handful of cells differ from paper Table 2 by more than 2 pts. All are explained: 1. **Prithvi v2 / Sentinel-2 tasks** — paper-era wrapper indexed S2 bands by HLS band names, silently feeding the model Red-Edge bands instead of NIR/SWIR. PR #385 / commit [`8746c342`](https://github.com/allenai/olmoearth_pretrain/commit/8746c342) fixed the mapping to `[B02, B03, B04, B08, B11, B12]` *after* the paper sweep was launched. **Our re-extracted Prithvi v2 embeddings use the corrected band mapping**, so they will give different (and more spectrally meaningful) downstream numbers than paper Table 2 reports. 2. **AnySat × m_forestnet** — paper-era code cropped m_forestnet's Landsat input to 320×320; that crop was reverted post-paper ([commit 96cf42b9](https://github.com/allenai/olmoearth_pretrain/commit/96cf42b9)). AnySat's wrapper picks `patch_size_meters` dynamically from the input height, so the larger image flips the patch grid. Other models on m_forestnet drift only ±1 pt because their wrappers are size-invariant. 3. **AnySat seg val/test (`m_cashew_plant`, `m_sa_crop_type`)** — AnySat returns *per-pixel* features; saving val/test at the full 256×256×1536 grid OOMs CPU RAM. Paper LP probes streamed val/test batch-by-batch and never had to hold them all. We instead save AnySat seg val/test at the same 1/16 subsample paper used for *training*. This is the only place our saved embeddings deviate from paper protocol — affects only AnySat's two largest seg tasks. **Pixel-level alignment caveat:** that 1/16 subsample uses an *unseeded* `randperm`, so the kept pixels within each AnySat seg image are random and not reproducible. Image (sample) order is still canonical, but for AnySat seg you cannot align *pixels* across models or to the label grid — treat each kept pixel's embedding+label pair as self-contained. All other (model, task) dumps are fully order-aligned per "Row ordering" above. 4. **AnySat H=W=1 time-series tasks** (`breizhcrops`, all `cropharvest_*`, `awf_*`, `nandi_*`) — paper-era AnySat ran fine on these; a later patch-size refactor unintentionally broke single-pixel inputs by choosing `patch_size > input_height`, crashing AnySat's LTAE `unfold(2, scale, scale)`. Fixed with the [PR #558](https://github.com/allenai/olmoearth_pretrain/pull/558) cap. Re-extracted embeddings reproduce paper Table 2 within typical LP variance (`breizhcrops`: 62.69 vs paper 62.7). 5. **`satlas / m_brick_kiln`** — paper Table 2 reports 83.0 but the paper-best W&B run (`oy47oe26`) actually scored 86.1. Our dump reproduces the W&B run (85.9), not the table. Looks like a transcription error in the table. A line-by-line table of paper-vs-extracted numbers (with all deviations explained) is in `verification.md`. ## Storage notes | Model class | Per-model size | |---|---| | Small ViT (Galileo nano/tiny, OlmoEarth nano) | ~10–25 GB | | ViT-B class | ~20–50 GB | | ViT-L / Huge class | ~30–85 GB | | ViT-7B (DINOv3) | ~52 GB each | | **AnySat** | **~575 GB** (per-pixel features for seg tasks) | `bfloat16` matches the autocast precision used at extraction time — casting to `float32` at load time loses no information but doubles memory. AnySat dominates the total because its segmentation outputs are per-pixel rather than per-patch (16× more vectors per image at input-resolution scale). ## Files in this dataset - `{model_group}/{task}/{train,valid,test}.pt` — embeddings + labels - `eval_settings/*.enriched.json` — paper-best per-task hyperparameters - `load_example.py` — minimal "load + run KNN" example - `verification.md` — full paper-vs-dump comparison table - `README.md` — this file ## Citation ```bibtex @article{olmoearth2025, title = {OlmoEarth: Stable Latent Image Modeling for Multimodal Earth Observation}, author = {Team OlmoEarth}, year = {2025}, eprint = {2511.13655}, archivePrefix = {arXiv} } ``` If you use the upstream foundation-model embeddings, please also cite the original models — see the Models table above. ## License CC BY 4.0. The underlying foundation models retain their original licenses; the embeddings here are derivatives of those models applied to public EO benchmarks (Geobench, BreizhCrops, CropHarvest, PASTIS, MADOS, Sen1Floods11) plus the paper's real-world AWF / Nandi datasets. ## Contact For questions or issues, open a GitHub issue at https://github.com/allenai/olmoearth_pretrain or email `[e-mail]`.