🤖 machine learning

Mini-JEPA Foundation Model Fleet Enables Agentic Hydrologic Intelligence

This paper introduces Mini-JEPAs, a fleet of specialized, lightweight foundation models that, when orchestrated by a routing agent, outperform a single large-scale geospatial model in hydrologic intelligence tasks while offering greater accessibility and computational efficiency.

Original authors: Mashrekur Rahman

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Mashrekur Rahman

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand the Earth's water systems, but you only have one giant, all-knowing encyclopedia. This encyclopedia, let's call it "AlphaEarth," is huge. It knows a little bit about everything: forests, deserts, rivers, and cities. It's great for general questions like, "What does this continent look like?"

But what if you ask a very specific question, like, "How wet is the soil right under this cloud?" or "How does the temperature change in this specific valley?" The giant encyclopedia might get a bit fuzzy here. It tried to learn everything at once, so it might have diluted the specific details needed for these tricky questions. Plus, this encyclopedia is so big and expensive that only a few people can even open it.

The researchers in this paper decided to try a different approach. Instead of one giant brain, they built a team of five small, specialized experts. They call this team the "Mini-JEPA Fleet."

The Team of Specialists

Think of these five Mini-JEPAs as five different detectives, each with a unique pair of glasses:

The Optical Detective: Wears glasses that see visible light (like a regular camera). It's great at seeing what plants are doing and what the ground looks like when the sun is out.
The Radar Detective: Wears glasses that see through clouds using radio waves. It can "feel" how rough the ground is or how wet the soil is, even when it's storming.
The Heat Detective: Wears thermal glasses that see temperature. It knows exactly how hot or cold the ground is, day or night.
The Seasonal Detective: Wears glasses that track how plants change over the four seasons. It understands the rhythm of spring growth and winter dormancy.
The Terrain Detective: Wears glasses that see the shape of the land and the type of soil. It knows the elevation and how water might flow down a hill.

How They Were Trained

The researchers didn't just guess what these experts should know. They trained them using a clever method called JEPA.

Usually, when you teach a computer to see, you show it a picture and ask it to redraw the missing parts pixel-by-pixel. That's like asking a student to copy a textbook word-for-word. It's slow and forces the student to memorize the "noise" (like static on a TV screen).

Instead, these Mini-JEPAs were taught to predict the meaning of the missing parts. It's like showing a detective a few clues and asking, "What is the likely story here?" This allowed them to learn the physics of their specific sensor (like how heat behaves or how radar bounces) without getting distracted by the messy details of the raw images.

Because they were all trained on the exact same computer architecture but with different "glasses," the researchers could prove that any differences in their brains came strictly from the sensors they used, not from how they were taught.

The "Router" Agent

Having five experts is great, but you need someone to decide who to ask. The researchers built a Router Agent (powered by a large language model).

When you ask a question like, "Is there flooding under these clouds?", the Router reads a tiny "ID card" for each expert. The ID card tells the Router what each expert is good at.

The Router sees "clouds" and "flooding."
It knows the Optical Detective can't see through clouds.
It knows the Radar Detective can see through clouds and is sensitive to water.
So, the Router sends the question specifically to the Radar Detective.

What They Found

The paper tested this system against the giant "AlphaEarth" encyclopedia and found some interesting things:

Specialization Works: Each Mini-JEPA was best at predicting the thing its sensor physically sees. The Heat Detective was amazing at predicting temperature; the Terrain Detective was amazing at predicting elevation.
Different Brain Shapes: Even though they are the same size, their "brains" (mathematically called manifolds) are shaped differently. The Heat Detective's brain is very simple and smooth (like a straight line), while the Radar Detective's brain is complex and jagged (like a crumpled piece of paper). This proves they are processing the world in fundamentally different ways.
Better Together: When the researchers combined the Mini-JEPAs with the giant AlphaEarth, they got better answers for specific things like soil moisture and rainfall than AlphaEarth could get alone.
The "Router" Wins: When the Router picked the right expert for the job, the answers were much more accurate and grounded in reality than when the giant encyclopedia tried to answer everything alone.

The Bottom Line

The paper argues that you don't always need a massive, expensive, planetary-scale supercomputer to understand the Earth. You can build a fleet of small, cheap, specialized models that can be trained on a single computer.

By using a smart "Router" to pick the right specialist for the job, you can get answers that are just as good (or sometimes better) than the giant models, but with much less computing power and more transparency. It's the difference between hiring one overworked generalist who knows a little about everything, versus hiring a team of five focused experts who know exactly what they are doing.

Technical Summary: Mini-JEPA Foundation Model Fleet Enables Agentic Hydrologic Intelligence

Problem Statement

Geospatial foundation models typically follow a "generalist" paradigm: training a single, massive model on diverse multispectral, radar, and atmospheric data to produce a universal embedding space. While effective for broad characterization, this approach faces three critical limitations in specialized domains like hydrology:

Signal Compromise: A single embedding space must allocate dimensions across phenomena with distinct physical origins (e.g., optical reflectance vs. microwave backscatter). This often leads to underweighting specialized signals relevant to specific hydrologic tasks.
Accessibility and Cost: Planetary-scale models (e.g., Google AlphaEarth) are expensive to train, frequently closed-source, and require massive compute resources, making them inaccessible for many research groups and difficult to audit or extend.
Lack of Specialization: Generalist models may not capture the nuanced physical signals required for specialized hydrologic reasoning, where the relevant data often resides in a specific sensor modality.

The authors propose that instead of a single monolithic model, a fleet of small, sensor-specialized models managed by an agentic router could offer a more efficient, interpretable, and performant alternative for hydrologic intelligence.

Methodology

1. The Mini-JEPA Fleet

The core contribution is a fleet of five Mini-JEPA (Joint Embedding Predictive Architecture) foundation models.

Architecture: All five models share an identical Vision Transformer-Small (ViT-S) backbone (12 layers, 6 attention heads, hidden dimension 384) and a 64-dimensional output space to ensure direct comparability with planetary-scale models like AlphaEarth.
Training Recipe: They utilize the I-JEPA training objective combined with a VICReg regularizer. Unlike Masked Autoencoders (MAE) that reconstruct pixels (forcing the model to learn sensor-specific noise), JEPA predicts masked latent representations. This allows the same training recipe to be applied across different sensor modalities without bias toward specific noise profiles (e.g., SAR speckle or MODIS striping).
Data: Each model is pretrained on a distinct Google Earth Engine satellite product covering 9,704 patches (128x128 pixels, 30m resolution) across the Continental United States (CONUS) for the year 2022. The five modalities are:
1. S2-Optical: Sentinel-2 surface reflectance (annual median).
2. S1-SAR: Sentinel-1 dual-polarization SAR (VV+VH).
3. MODIS-Thermal: MODIS land surface temperature (day/night).
4. S2-Phenology: Sentinel-2 quarterly composites capturing seasonal dynamics.
5. Topo-Soil: A stack of SRTM elevation/slope/aspect and SoilGrids properties.

2. Evaluation Framework

The authors characterize the fleet along three axes:

Dimension-Level Interpretability: Using Spearman correlations and Random Forest permutation importance to map embedding dimensions to environmental variables (e.g., soil moisture, precipitation, temperature).
Manifold Geometry: Analyzing the geometric structure of the 64-d embedding spaces via global participation ratios (effective dimensionality) and local intrinsic dimensionality (k-NN based).
Complementarity: Using Canonical Correlation Analysis (CCA) and joint Random Forest regression to determine if Mini-JEPAs provide information not already captured by AlphaEarth.

3. Agentic System

A routing agent (LLM) was implemented to operationalize the fleet:

Reference Cards: Each Mini-JEPA is summarized in a compact "reference card" containing its dimension dictionary, geometric profile, physical signal description, and predictive performance metrics.
Routing Logic: Given a natural language hydrologic question, the router LLM reads these cards and emits a structured plan to query specific Mini-JEPAs (and optionally AlphaEarth).
Retrieval & Synthesis: Selected models perform k-nearest-neighbor retrieval against FAISS indices. A synthesis LLM combines the retrieved context to generate a final answer.

Key Results

1. Sensor Physics Drives Specialization

Each Mini-JEPA demonstrated its strongest predictive skill on the environmental variable physically matched to its sensor:

MODIS-Thermal: Achieved $R^2 = 0.97$ for mean annual temperature.
Topo-Soil: Achieved $R^2 = 0.97$ for elevation.
S2-Phenology: Achieved $R^2 = 0.81$ for precipitation (via vegetation response).
S1-SAR: Achieved $R^2 = 0.62$ for precipitation.
S2-Optical: Achieved $R^2 = 0.73$ for aridity.
Notably, the fleet's performance patterns were not simple sensor-to-variable lookups but reflected complex surface processes (e.g., phenology predicting precipitation better than SAR). All models struggled with categorical land cover (NLCD), suggesting a task limitation rather than a model failure.

2. Distinct Manifold Geometries

The five models produced embedding manifolds with distinct geometric structures, reflecting their sensor physics:

Global Dimensionality: Ranged from 8.9 (S1-SAR) to 20.2 (MODIS-Thermal).
Local Complexity: S1-SAR exhibited high local complexity (mean local $n_{80} \approx 6.0$ ), while MODIS-Thermal was nearly one-dimensional locally ( $n_{80} \approx 2.0$ ), consistent with temperature varying along smooth gradients.
This confirms that fixing the architecture and recipe isolates the sensor as the primary driver of structural differences.

3. Complementarity with AlphaEarth

Additive Value: Combining Topo-Soil or S2-Phenology with AlphaEarth yielded positive gains ( $\Delta R^2$ up to 0.031) for soil moisture, aridity, and precipitation.
Redundancy: For variables where AlphaEarth already performed well (temperature, elevation), the Mini-JEPAs offered little additional gain, though they could match AlphaEarth's performance individually.
Substitution: On physics-matched variables, a single Mini-JEPA (22M parameters) could substitute for the planetary-scale generalist, achieving comparable $R^2$ scores with a fraction of the compute.

4. Agentic Performance

Routing Accuracy: The router LLM achieved a 100% hit rate on expected modalities for a curated set of 40 questions.
Response Quality: In "LLM-as-Judge" evaluations:
- Aggregate: The dual-retrieval system (AlphaEarth + Fleet) showed only modest aggregate gains over AlphaEarth alone ( $\Delta \approx +0.021$ ), largely due to score saturation on broad questions.
- Single-Modality Questions: For questions where the physical signal sits cleanly in one sensor, the routed fleet significantly outperformed AlphaEarth alone (Cohen's $d = 1.10, p = 0.031$ ).
- Substitution: The routed fleet alone performed comparably to AlphaEarth on physics-matched questions, validating its use as a standalone alternative when generalists are unavailable.

Significance and Claims

The paper claims that locally-trained, specialized Mini-JEPAs offer a viable, complementary path to planetary-scale foundation models for environmental science.

Democratization: The fleet design lowers the barrier to entry, allowing institutions with modest compute resources (a single workstation GPU) to build, audit, and extend bespoke environmental embeddings.
Interpretability: By separating models by sensor physics, the system provides clearer interpretability regarding which physical signals drive specific predictions, unlike the "black box" nature of generalist embeddings.
Agentic Flexibility: The routing agent successfully composes these specialists into a functional system that adapts to the physical requirements of a query, outperforming generalists on specialized tasks.
Coexistence: The authors envision a future where generalist embeddings and specialist fleets coexist, with agents dynamically selecting the appropriate tool based on the physics of the question.

The authors remain modest, acknowledging limitations such as the restriction to CONUS data, the single-year training corpus, and the saturation of LLM-as-Judge scores on broad questions. They suggest future work should explore specialization by hydrologic regime (e.g., snow, arid) and temporal scale, and test geographic generalization beyond the training domain.