UniTeD: Resolving Sub-Grid Heterogeneity for Robust Grid Cell Prediction

Machine learning (ML) has seen widespread adoption in the physical sciences for developing data-driven emulators of grid-based simulations. Despite this rapid progress, physical systems characterized by sub-grid heterogeneity remain largely underexplored. In domains ranging from land surface modeling in Earth sciences to multi-scale stress analysis in materials science, a grid cell's state is far from uniform. Instead, the cell's state is highly heterogeneous: distinct sub-grid “tiles” exhibit intrinsic processes that must be resolved before computing the grid cell's aggregate behavior. Directly applying standard ML models, which lack the inductive bias to account for this heterogeneity or resolve these underlying processes, can lead to structural entanglement and poor generalization when sub-grid fractions shift. To address this gap, we introduce UniTeD (Unified Tiled-Expert Demixer), a sub-grid-aware ML architecture that strictly enforces physical aggregation logic. UniTeD isolates latent constitutive functions purely from aggregate supervision, enabling robust generalization across varying sub-grid fraction compositions. Evaluated on synthetic benchmarks and global ERA5 surface energy flux inference, UniTeD achieves zero-shot compositional generalization, reducing prediction errors by up to 28% relative to the strongest baseline.