# Attention U-Net as a surrogate model for groundwater prediction

arXiv: https://arxiv.org/abs/2204.04518 · HTML: https://mlsanity.com/paper/2204.04518/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 47 references that are in the map.
- First posted on arXiv 9 Apr 2022.

## Builds on (most cited first, 18 of 18)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Image-to-Image Translation with Conditional Adversarial Networks](https://mlsanity.com/paper/1611.07004/) (2016)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Going Deeper with Convolutions](https://mlsanity.com/paper/1409.4842/) (2014)
- [Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translation](https://mlsanity.com/paper/1609.08144/) (2016)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [Depth Map Prediction from a Single Image using a Multi-Scale Deep Network](https://mlsanity.com/paper/1406.2283/) (2014)
- [Transformers in Vision: A Survey](https://mlsanity.com/paper/2101.01169/) (2021)
- [Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks](https://mlsanity.com/paper/1602.07868/) (2016)
- [Attention Gated Networks: Learning to Leverage Salient Regions in Medical Images](https://mlsanity.com/paper/1808.08114/) (2018)
- [Bayesian Deep Convolutional Encoder-Decoder Networks for Surrogate Modeling and Uncertainty Quantification](https://mlsanity.com/paper/1801.06879/) (2018)
- [Modeling the Dynamics of PDE Systems with Physics-Constrained Deep Auto-Regressive Networks](https://mlsanity.com/paper/1906.05747/) (2019)
- [An Attentive Survey of Attention Models](https://mlsanity.com/paper/1904.02874/) (2019)
- [Attention in Natural Language Processing](https://mlsanity.com/paper/1902.02181/) (2019)
- [Deep convolutional encoder-decoder networks for uncertainty quantification of dynamic multiphase flow in heterogeneous media](https://mlsanity.com/paper/1807.00882/) (2018)
- [Structured Bayesian Gaussian process latent variable model: applications to data-driven dimensionality reduction and high-dimensional inversion](https://mlsanity.com/paper/1807.04302/) (2018)
- [Discovering state-parameter mappings in subsurface models using generative adversarial networks](https://mlsanity.com/paper/1810.12856/) (2018)

## Built on by (most cited first, 3 of 3)

- [AI enhanced data assimilation and uncertainty quantification applied to Geological Carbon Storage](https://mlsanity.com/paper/2402.06110/) (2024)
- [Auto-Regressive U-Net for Full-Field Prediction of Shrinkage-Induced Damage in Concrete](https://mlsanity.com/paper/2509.20507/) (2025)
- [Hybrid Lagrangian-Eulerian Model for Lagrangian Fluid Simulation](https://mlsanity.com/paper/2608.01164/) (2026)

## Abstract (arXiv)

Numerical simulations of groundwater flow are used to analyze and predict the response of an aquifer system to its change in state by approximating the solution of the fundamental groundwater physical equations. The most used and classical methodologies, such as Finite Difference (FD) and Finite Element (FE) Methods, use iterative solvers which are associated with high computational cost. This study proposes a physics-based convolutional encoder-decoder neural network as a surrogate model to quickly calculate the response of the groundwater system. Holding strong promise in cross-domain mappings, encoder-decoder networks are applicable for learning complex input-output mappings of physical systems. This manuscript presents an Attention U-Net model that attempts to capture the fundamental input-output relations of the groundwater system and generates solutions of hydraulic head in the whole domain given a set of physical parameters and boundary conditions. The model accurately predicts the steady state response of a highly heterogeneous groundwater system given the locations and piezometric head of up to 3 wells as input. The network learns to pay attention only in the relevant parts of the domain and the generated hydraulic head field corresponds to the target samples in great detail. Even relative to coarse finite difference approximations the proposed model is shown to be significantly faster than a comparative state-of-the-art numerical solver, thus providing a base for further development of the presented networks as surrogate models for groundwater prediction.