🤖 machine learning

OceanCBM: A Concept Bottleneck Model for Mechanistic Interpretability in Ocean Forecasting

The paper introduces OceanCBM, the first concept bottleneck model for ocean forecasting that achieves interpretable, physically grounded predictions of marine heatwave precursors by routing information through a mix of prescribed geophysical concepts and a learned free concept, thereby ensuring mechanistic consistency without sacrificing predictive skill.

Original authors: Sanah Suri, Kieran Ringel, Maike Sonnewald

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Sanah Suri, Kieran Ringel, Maike Sonnewald

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 predict a massive storm at sea. A standard computer model might be like a super-smart weather forecaster who can tell you, with 99% accuracy, that a storm is coming. However, if you ask why it's coming, the model just shrugs. It sees patterns in the data that humans can't understand, making it a "black box." If the weather changes in a way the model hasn't seen before, it might fail spectacularly because it didn't actually understand the physics of the storm; it just memorized the symptoms.

The paper "OceanCBM" introduces a new kind of AI model designed to fix this problem. Think of it as giving the computer a checklist of physical rules it must use to make its prediction, while still letting it figure out the details on its own.

Here is how it works, broken down into simple concepts:

1. The Problem: The "Black Box" vs. The "Black Hole"

Ocean scientists need to predict Marine Heatwaves (periods where the ocean gets dangerously hot).

Old AI models are like a magician pulling a rabbit out of a hat. You see the rabbit (the prediction), but you don't know how the trick was done.
Strict Physics models are like a rigid recipe book. They follow the rules perfectly but might fail if the ingredients (the ocean conditions) are slightly different from what the recipe expects.

2. The Solution: The "Concept Bottleneck"

The authors created OceanCBM. Imagine the AI's brain has a narrow hallway (a bottleneck) that all information must pass through before making a final prediction.

Instead of letting the AI guess freely, the researchers put four specific "checklist items" (concepts) in that hallway that are based on real ocean physics:

Heat Exchange: How much heat is moving between the air and the water.
Mixing Speed: How fast the water is churning vertically.
Water Density: How much the water resists mixing (like oil on water).
Depth Changes: How deep the warm layer of water is getting.

The AI is forced to look at the ocean data and answer: "Based on these four rules, what is the heat content?"

3. The Secret Ingredient: The "Free Concept"

Here is the clever twist. The researchers knew their checklist wasn't perfect. The ocean is complex, and they might have missed a rule. So, they added a fifth, "free" concept.

Think of this as a "Catch-All" bucket.

If the four physical rules explain everything, the bucket stays empty.
If there is something weird happening that the rules don't cover (like a hidden current or a new type of mixing), the AI puts that leftover information into the "Free Concept" bucket.

This prevents the AI from getting stuck trying to force a square peg into a round hole. It allows the model to be flexible while still being grounded in science.

4. The "Team of Experts" (Ensemble)

To make sure the model isn't just lucky, they trained five different versions of this AI (an ensemble).

The Test: They asked, "Do all five experts agree on why the heatwave is happening?"
The Result:
- Models that only looked at the final answer (no checklist) gave totally different, chaotic reasons for the same prediction.
- Models that only followed the rules (no "free" bucket) were inconsistent.
- OceanCBM (the mix of rules + free bucket) showed that all five experts agreed on the main physical reasons. They all pointed to the same "drivers" of the heatwave.

5. Real-World Proof: The 2012 Gulf of Maine Heatwave

The team tested their model on a real event: a massive heatwave in the Gulf of Maine in 2012.

What the model found: It correctly identified that the heatwave started because the air dumped a lot of heat into the ocean, and then the ocean water became "stiff" (stratified), trapping that heat near the surface like a lid on a pot.
Why it matters: This matched what human scientists already knew, but it also showed how the heat was trapped. It confirmed that the model wasn't just guessing; it was actually "thinking" like a physicist.

The Bottom Line

OceanCBM is a new type of AI that balances accuracy with understanding.

It predicts ocean heatwaves just as well as the best "black box" models.
But unlike them, it explains its reasoning using real physical concepts.
It uses a "safety net" (the free concept) to catch anything the known rules miss, ensuring the model stays accurate even when the ocean does something unexpected.

In short, it's a model that doesn't just tell you what will happen, but helps you understand why it's happening, making it a more trustworthy tool for scientists studying our changing climate.

Technical Summary: OceanCBM

Problem Statement

Predicting extreme ocean phenomena, such as Marine Heatwaves (MHWs), presents a dual challenge: achieving high predictive accuracy while maintaining mechanistic interpretability. While recent machine learning (ML) models demonstrate strong forecasting skills for ocean states, they often function as "black boxes," learning spurious correlations rather than underlying physical drivers. Existing approaches to embedding physics into ML fall into two categories with distinct limitations:

Physics-Informed Neural Networks (PINNs): These enforce dynamical constraints directly but can be too rigid, failing to generalize to regimes outside historical records or converging to physically incorrect spurious solutions.
Post-hoc Explainability (XAI): These methods analyze trained models to determine input reliance but explain what the model did rather than ensuring the model adhered to physical laws during training.

Consequently, there is a lack of frameworks that balance predictive skill with guaranteed physical fidelity, particularly for spatiotemporal ocean dynamics where understanding the drivers of MHWs (which occur in the mixed layer, not just at the surface) is critical for ecological and socioeconomic forecasting.

Methodology: OceanCBM

The authors introduce OceanCBM, the first Concept Bottleneck Model (CBM) designed for spatiotemporal ocean physics prediction. The model is designed to forecast Mixed Layer Heat Content (MLHC), a key precursor to MHWs, by routing predictions through an intermediate layer of interpretable concepts.

Data and Input

Dataset: The model utilizes the Ocean Reanalysis System 5 (ORAS5), which provides monthly averaged ocean fields (1979–2018) at 1/4° resolution. Uniquely, the authors leverage all five ensemble members of ORAS5 as independent training samples to capture internal variability and data uncertainty, rather than relying on a single control member or ensemble mean.
Input: The model processes a 6-month context window ( $t-6$ to $t-1$ ) of 12 spatially collapsed 2D physical variables (e.g., sea surface height, zonal/meridional velocity) over the North Atlantic basin.
Preprocessing: Data is detrended to focus on climatological variability rather than long-term warming trends.

Architecture

OceanCBM employs an ensemble of five U-Net-based architectures. The workflow is as follows:

Encoder: A standard U-Net encoder with 4 downsampling stages (channels: 64, 128, 256, 512) using convolutions, batch normalization, ReLU, and max pooling.
Concept Bottleneck: The encoder outputs pass through a bottleneck layer containing two types of concepts:
- Prescribed Concepts: Four fields derived from geophysical fluid dynamics: (1) Heat flux entrainment, (2) Vertical shear, (3) Buoyancy frequency, and (4) Mixed Layer Depth (MLD) tendency.
- Free Concept: A single unconstrained field designed to capture residual physical processes not represented by the prescribed set.
Decoder: A U-Net decoder reconstructs the spatial fields.
Output Heads: The model produces two outputs: the predicted concept fields (prescribed + free) and the final target (MLHC). The final prediction of MLHC is a linear combination of the concept representations.

Training Strategy

The model is trained using mixed supervision with a weighted loss function:
$L_{combined} = \lambda L_{concept} + (1 - \lambda)L_{pred}$

Adaptive $\lambda$ Schedule: Training begins with $\lambda = 0.8$ to prioritize learning the physical concepts early, then decays exponentially to $\lambda = 0.2$ over 100 epochs. This ensures the bottleneck establishes physically meaningful representations while maintaining a degree of concept supervision throughout training.
Ensemble: Five models are trained with independent random weight initializations to assess robustness and consistency in learned mechanistic pathways.

Key Contributions

First Spatiotemporal CBM for Ocean Physics: OceanCBM extends the CBM framework to complex, continuous spatiotemporal data, enabling interpretable forecasting of MLHC and mechanistic interrogation of MHW drivers.
Hybrid Supervision with a Free Concept: The authors propose a novel bottleneck architecture that combines prescribed physical concepts with a "free" concept. This free concept serves a dual role:
- It acts as a regularizer for the prescribed concepts.
- It captures residual physical processes and emergent phenomena not covered by the prescribed set, preventing the model from being over-constrained.
Demonstration of Consistent Mechanistic Learning: The study shows that mixed supervision yields consistent mechanistic representations across ensemble members, whereas baselines (prediction-only or prescription-only) learn highly variable latent structures despite similar predictive performance.
Real-World Case Study: The framework is applied to the 2012 Gulf of Maine MHW, successfully recovering known physical drivers and providing mechanistic insights into the event's onset.

Results

Predictive Skill

Accuracy: OceanCBM achieves an Anomaly Correlation Coefficient (ACC) of 0.91–0.92 for MLHC prediction across seasons, which is statistically indistinguishable from a "prediction-only" baseline (no concept bottleneck) and a "prescription-only" baseline. This indicates that introducing interpretability constraints does not sacrifice predictive skill.
Concept Prediction: The model successfully predicts the prescribed concepts (e.g., heat flux entrainment ACC of 0.77, buoyancy frequency ACC of 0.81 in summer), demonstrating that the model learns meaningful representations rather than weak approximations.

Mechanistic Consistency

Ensemble Variability: When analyzing the linear weights connecting concepts to the final MLHC prediction:
- Prediction-only models show high variability in latent representations across ensemble members, making it impossible to discern a consistent learning strategy.
- Prescription-only models tend to latch onto specific concepts (e.g., vertical shear) inconsistently.
- OceanCBM exhibits significantly more consistency in how it weights the prescribed and free concepts across different random initializations, suggesting the mixed supervision approach guides the model toward stable, physically meaningful pathways.

Role of the Free Concept

Residual Channel: The free concept captures predictive structure in dynamically active regions (e.g., Gulf Stream, Labrador Current) where prescribed concepts diverge from the target.
Regularization: Empirical results show that the presence of the free concept improves the prediction of difficult concepts like vertical shear (e.g., ACC increases from 0.53 to 0.60 in winter for OceanCBM vs. prescription-only), acting as a structured regularizer that absorbs residual variability.

Case Study: 2012 Gulf of Maine MHW

A retrospective analysis of the 2012 event revealed:

Onset Mechanism: The model identified a positive anomaly in MLHC and buoyancy frequency (indicating enhanced stratification) prior to the event, which likely trapped heat near the surface.
Drivers: The analysis confirmed that the event was driven largely by atmospheric heat flux anomalies and heat retention, consistent with prior literature, while also highlighting the role of subsurface processes in sustaining the heat.
Ecological Alignment: The predicted warm mixed layer conditions aligned with documented ecological impacts, such as increased lobster landings and declines in northern shrimp.

Significance and Claims

The paper claims that OceanCBM successfully addresses the trade-off between predictive performance and interpretability in high-stakes scientific domains. By integrating mixed supervision with a free concept, the framework:

Balances Constraint and Flexibility: It guides the model toward known physics without rigidly over-constraining it, allowing for the discovery of unknown or unresolved mechanisms.
Ensures Reliability: It provides a pathway for models to learn consistent mechanistic drivers, reducing the risk of models predicting correctly for the wrong reasons (spurious correlations).
Offers a New Paradigm: It demonstrates that "soft" physical structure, combined with ensemble methods, can yield interpretable representations that are robust to initialization and suitable for diagnosing extreme events in a changing climate.

The authors conclude that this approach offers a grounded pathway for learning interpretable representations of complex dynamical processes, with potential transferability to broader Earth system models.