OceanMoE: Structured Conditional Sparse Computation for Long-Horizon Multivariate Ocean Forecasting
Yishun Zhu, Jian Wang
cs.LG
2026-09-17
OceanMoE routes sparse experts by target and grid cell. On 60-month ORAS5 rollouts, geometric-mean RMSE falls 4.91% vs reproduced ORCA-DL and 8.34% vs the 10-variable baseline.
Multivariate ocean forecasting needs two things that fight each other. Temperature, salinity, currents, sea-surface height, and fluxes share one coupled circulation, so a single model can keep that context. Those fields still differ in spatial structure, timescale, magnitude, and difficulty, and the same field can need different computation in different basins.
A fully shared network reuses one parameter set everywhere. Coupling survives; specialization does not. Training a separate model per variable buys specialization and breaks the shared ocean state. Data-driven Earth-system models such as Pangu-Weather, GraphCast, and ORCA-DL already commit to unified representations. OceanMoE asks the next design question: what should stay shared, and where should computation branch by target and location.
The backbone still follows ORCA-DL: variable-wise encoders, a fusion stage, variable-wise decoders. One forward pass predicts a six-month block. The last predicted ocean state seeds the next block, for ten blocks and 60 months. Surface wind stresses τu and τv are prescribed external conditions. The ocean model does not forecast them.
The fusion stage is where the architecture changes. Five pieces matter:
The training loss is latitude-area-weighted normalized RMSE, minus a Jensen–Shannon term (coefficient 10^{-3}) between target-averaged routing distributions, which pushes different targets off the same experts. The pool has E=14 experts. Training uses monthly historical runs from 20 CMIP6 models for 1850–2009. SODA2 is the training-time validation archive. The locked benchmark is ORAS5: every model starts from the same December 2020 ocean state and rolls through December 2025. Fields sit on a 1° 128×360 grid with 16 vertical levels down to 1000 m.
Parameter counts match the comparison contracts. Reproduced ORCA-DL is 540.5M on six ocean variables. ORCA-DL-Expanded is 883.4M on the 14-variable contract. OceanMoE is 936.2M: 37.2M in private experts, 5.5M in routing and target conditioning, 13.3M in shared pathways, and the rest in the variable-wise encoder–decoder backbone. Sparse dispatch evaluates only selected token–expert pairs.
The headline score is a hierarchical geometric mean of physical-unit RMSE ratios, with equal weight per variable and equal weight per vertical level. Negative values mean OceanMoE has lower error.
| Setting | Baseline | Geometric relative RMSE drop |
| Six variables | reproduced ORCA-DL | 4.91% |
| Ten variables | ORCA-DL-Expanded | 8.34% |
| Ten variables without sea-surface salinity | same, remaining nine | 5.91% |
In physical units on the six-variable task, salinity so falls from 0.206 to 0.170 psu, potential temperature thetao from 1.048 to 0.928 °C, and sea-surface height zos from 0.093 to 0.081 m, about 12.9%. Sea-surface temperature tos, zonal current uo, and meridional current vo get worse. On the ten-variable task, sea-surface salinity sos falls from 0.677 to 0.490 psu and mixed-layer thickness mlotst from 41.4 to 36.4 m; downward heat flux hfds rises from 80.0 to 83.7 W m^{-2}.
Regionally, six of seven diagnostic boxes improve on the six-variable task. The northern high-latitude band and the equatorial Pacific drop 14.03% and 13.95%; the Southern Ocean is 4.66% worse. All seven boxes improve on the ten-variable task, with 19.06% in the northern high-latitude band and 18.47% in the Southern Ocean. Over forecast months 13–60, OceanMoE sits below the matching baseline in all 48 months of both contracts. The first 12 months still cross.
Component-removal tests on the ten-variable score, relative to the full model: drop shared pathways +1.20%, drop target-specific fusion +1.84%, drop the decoder geographic router +1.92%. Replace dynamic selection with fixed top-2 and the geometric ratio jumps +17.69%. No stripped variant beats the full model on the variable-balanced score.
Routing diagnostics line up with the design. The main run uses 1.76 encoder experts and 2.15 decoder experts per position on average, while the effective expert counts across targets, locations, and months reach 3.22 and 6.93. A spherical-harmonic fit to the mean decoder routing map reaches a weighted R² of about 0.935. Sea-surface salinity is more extreme: encoder K=1.03, decoder K=2.63. In the fusion heatmap, velocity targets keep large self-weights (uo 0.659, vo 0.632); tos mixes thetao, tos, uo, and vo to a combined 0.883. The paper treats those weights as consistent with known ocean structure, not as causal claims.
Earth-system forecast groups have been stuck between one shared network and a pile of per-variable models. OceanMoE is a workable middle: keep a shared representation, branch sparse computation by target and grid cell, and park the context that should not be partitioned on residual and seasonal paths.
What transfers is not the specific 14-expert pool. It is target-conditioned fusion into content before routing, a spherical-harmonic geographic prior on the decoder, and a confidence threshold that sets expert count. Fixed top-2 is the ablation that hurts most, so on these ocean tokens the number of active experts is itself a capacity knob.
The gain is incremental. The baselines are reproduced ORCA-DL variants, not eddy-resolving systems such as XiHe or WenHai. The grid is 1°, far from operational medium-range resolution. The setting that matches the paper is seasonal-to-interannual multivariate rollout, not next-day mesoscale eddies.
The six-variable comparison is a cross-contract score. OceanMoE is a 14-variable checkpoint evaluated on the six ORAS5 targets that have complete series; reproduced ORCA-DL was trained on the six-variable contract. Parameter counts are 936.2M versus 540.5M, and the selected checkpoints differ (12k / 16k / 28k). Geometric averaging hides regressions on tos, uo, and vo. The Southern Ocean gets worse in the six-variable setting. The high-latitude boxes are 60°–63.5° boundary strips, not full polar basins. Atmospheric stresses are prescribed; the model does not roll the atmosphere. Training is CMIP6 historical simulation and evaluation is ORAS5 for 2021–2025, with no separate split of reanalysis mismatch. There is no independent compute-cost table, only the statement that sparse dispatch skips unselected pairs.