K-STEMIT: Knowledge-Informed Spatio-Temporal Efficient Multi-Branch Graph Neural Network for Subsurface Stratigraphy Thickness Estimation from Radar Data

Zesheng Liu; Maryam Rahnemoonfar

cs.LGApr 10, 2026

K-STEMIT: Knowledge-Informed Spatio-Temporal Efficient Multi-Branch Graph Neural Network for Subsurface Stratigraphy Thickness Estimation from Radar Data

Authors: Zesheng Liu, Maryam Rahnemoonfar

Organizations: Department of Computer Science and Engineering, Lehigh University, 113 Research Drive, Bethlehem, 18015, Pennsylvania, USA · Department of Civil and Environmental Engineering, Lehigh University, 19 Memorial Drive W., Bethlehem, 18015, Pennsylvania, USA

Abstract

Subsurface stratigraphy contains important spatio-temporal information about accumulation, deformation, and layer formation in polar ice sheets. In particular, variations in internal ice layer thickness provide valuable constraints for snow mass balance estimation and projections of ice sheet change. Although radar sensors can capture these layered structures as depth-resolved radargrams, convolutional neural networks applied directly to radar images are often sensitive to speckle noise and acquisition artifacts. In addition, purely data-driven methods may underuse physical knowledge, leading to unrealistic thickness estimates under spatial or temporal extrapolation. To address these challenges, we develop K-STEMIT, a novel knowledge-informed, efficient, multi-branch spatio-temporal graph neural network that combines a geometric framework for spatial learning with temporal convolution to capture temporal dynamics, and incorporates physical data synchronized from the Model Atmospheric Regional physical weather model. An adaptive feature fusion strategy is employed to dynamically combine features learned from different branches. Extensive experiments have been conducted to compare K-STEMIT against current state-of-the-art methods in both knowledge-informed and non-knowledge-informed settings, as well as other existing methods. Results show that K-STEMIT consistently achieves the highest accuracy while maintaining near-optimal efficiency. Most notably, incorporating adaptive feature fusion and physical priors reduces the root mean-squared error by 21.01% with negligible additional cost compared to its conventional multi-branch variants. Additionally, our proposed K-STEMIT achieves consistently lower per-year relative MAE, enabling reliable, continuous spatiotemporal assessment of snow accumulation variability across large spatial regions.

Figures & tables

Figure 1: Diagram of capturing radargrams and generating corresponding labels. ( a ) How airborne radar sensor is used to capture the status of internal ice layers (Image adapted from [ 40 ] ) ( b ) Example of a radargram captured by the airborne radar sensor. ( c ) Corresponding labeled images, where each ice layer is manually annotated. Each layer is formed in a certain year.

Figure 2: Architecture of our proposed K-STEMIT.

Figure 3: Dimension reduction in the spatial branch, where the graph sequence is compressed into a single graph with concatenated node features.

Stage	Input shape	Output shape	Description
Input graph sequence	$\mathbf{X}\in\mathbb{R}^{m\times 256\times 8}$	$\mathbf{X}\in\mathbb{R}^{m\times 256\times 8}$	Sequence of $m$ observed spatial graphs. Each graph contains 256 nodes, and each node has 8 input features.
Spatial branch: dimension reduction	$\mathbb{R}^{m\times 256\times 8}$	$\tilde{\mathbf{X}}_{\text{spatial}}\in\mathbb{R}^{256\times(6m+2)}$	Rearranges and compresses the multi-layer node features into a single graph representation for spatial modeling.
Spatial branch: GraphSAGE layer	$\mathbb{R}^{256\times(6m+2)}$	$\mathbf{h}_{\text{spatial}}\in\mathbb{R}^{256\times d^{\prime}}$	Learns localized spatial dependencies over the graph topology from the compressed node features.
Temporal branch: dimension reduction	$\mathbb{R}^{m\times 256\times 8}$	$\tilde{\mathbf{X}}_{\text{temporal}}\in\mathbb{R}^{m\times 256\times 6}$	Projects each observed layer into a reduced node-wise feature representation for temporal modeling.
Temporal branch: temporal convolution	$\mathbb{R}^{256\times m\times 6}$	$\mathbf{h}_{\text{temporal}}\in\mathbb{R}^{256\times d^{\prime}}$	Extracts cross-layer temporal dependencies for each node across the $m$ observed layers.
Fusion (element-wise addition)	$\mathbf{h}_{\text{spatial}},\mathbf{h}_{\text{temporal}}\in\mathbb{R}^{256\times d^{\prime}}$	$\mathbf{h}\in\mathbb{R}^{256\times d^{\prime}}$	Fuses the spatial and temporal branch outputs; both branches share the same fusion dimension $d^{\prime}$ .

Table 1: Stage-wise summary of tensor dimensions in K-STEMIT for the experimental setting used in this study. Batch dimension is omitted for clarity. Here,

m

denotes the number of observed upper ice layers and

n

denotes the number of target deeper layers. Hidden dimensions are kept symbolic where appropriate to preserve generality.

Figure 4: Diagram of Temporal Convolution Block.

Learning rate	Weight decay	Scheduler	Scheduler arg	$\eta_{\min}$	Validation RMSE
$10^{-2}$	$10^{-5}$	Cosine	450	$10^{-7}$	2.5620 $\pm$ 0.0571
$10^{-3}$	$10^{-5}$	Cosine	450	$10^{-7}$	2.7744 $\pm$ 0.0691
$5\times 10^{-3}$	$10^{-4}$	Cosine	450	$10^{-7}$	2.4382 $\pm$ 0.0917
$5\times 10^{-3}$	$10^{-6}$	Cosine	450	$10^{-7}$	2.4137 $\pm$ 0.0730
$5\times 10^{-3}$	$10^{-5}$	Step	50	–	2.6098 $\pm$ 0.0902
$5\times 10^{-3}$	$10^{-5}$	Step	75	–	2.5148 $\pm$ 0.0656

Table 2: Hyperparameter sensitivity study for K-STEMIT. Each case varies one component around the final configuration unless otherwise noted. Validation RMSE is reported for model selection. For the step scheduler, “Scheduler arg” denotes the decay interval in epochs; for cosine annealing, it denotes

T_{\max}

. The final selected configuration is highlighted in bold.

Figure 5: Qualitative results of our proposed K-STEMIT. The blue line is used to generate the graphs. The green line is the groundtruth (manually-labeled ice layers) and the red line is the model prediction.

Model	Physical Node Feature?	RMSE	MAE	Average Training Time (Seconds)
GCN [ 27 ]	x	5.0573 $\pm$ 0.0874	3.7655 $\pm$ 0.0509	1099
GIN [ 71 ]	x	3.5696 $\pm$ 0.0591	2.5416 $\pm$ 0.0521	862
AGCN-LSTM [ 85 ]	x	3.4808 $\pm$ 0.0397	2.4486 $\pm$ 0.0438	9404
GCN-LSTM [ 86 ]	x	3.1745 $\pm$ 0.1045	2.2400 $\pm$ 0.0923	7441
GraphSAGE-LSTM [ 31 ]	x	3.3837 $\pm$ 0.1103	2.4380 $\pm$ 0.1171	4579
GraphSAGE [ 16 ]	x	3.1873 $\pm$ 0.0588	2.1831 $\pm$ 0.0233	842

Table 3: Experiment results of K-STEMIT, its different variants, current state-of-the-art models and other baselines. Results are reported as the mean and standard deviation of the RMSE and MAE on the test dataset over five individual trials. Train time is reported as the average train time over five individual trials for 450 epochs.

Snow Mass Balance	Average Yearly Surface Temperature	Height Change Due To Refreezing	Height Change Due To Melting	Snowpack Height	Average RMSE	Improvement(%)
					3.1170 $\pm$ 0.0780	0.00
				✓	2.4647 $\pm$ 0.0588	20.93
			✓		2.8627 $\pm$ 0.0798	8.15
			✓	✓	2.4603 $\pm$ 0.0657	21.07
		✓			2.8877 $\pm$ 0.0621	7.36
		✓		✓	2.4831 $\pm$ 0.0544	20.34

Table 4: Ablation study on the choice of different physical features

Model	Results for STEMIT	Results for K-STEMIT
GCN+GraphSAGE+TempConv (Three-branch, No Clamp)	3.1204 $\pm$ 0.0801	2.4468 $\pm$ 0.0777
GCN+GraphSAGE+TempConv (Three-branch, Clamp)	3.1089 $\pm$ 0.0569	2.4545 $\pm$ 0.0707
GCN+TempConv (Spectral + Temporal)	3.2461 $\pm$ 0.0529	2.5190 $\pm$ 0.0529
GCN+GraphSAGE (Spectral + Spatial)	3.2195 $\pm$ 0.1183	2.8079 $\pm$ 0.0709
GCN (Spectral)	5.0573 $\pm$ 0.0874	3.3298 $\pm$ 0.0491
GraphSAGE (Spatial)	3.1873 $\pm$ 0.0588	2.6173 $\pm$ 0.0835

Table 5: Experiment results of ablation study on the combination of different branches for STEMIT and K-STEMIT. STEMIT denotes the non-knowledge-informed version of K-STEMIT that removes the MAR physical node features. Results are reported as the mean and standard deviation of the RMSE on the test dataset over five individual trials.

Model	Multi-Branch Architecture	Adaptive Feature Fusion Strategy	Knowledge Informed Learning	Average RMSE	Improvement (%)
Multi-Branch	✓			3.0980 $\pm$ 0.0857	0.00
STEMIT	✓	✓		3.1170 $\pm$ 0.0780	-0.61
Knowledge-Informed Multi-Branch	✓		✓	2.4892 $\pm$ 0.0556	19.65
K-STEMIT	✓	✓	✓	2.4168 $\pm$ 0.0553	21.99

Table 6: RMSE values and percentage improvements for different variants of K-STEMIT.

Fusion Strategy	Fusion Coefficient	RMSE	MAE	Time
Global adaptive fusion, used in K-STEMIT	$\alpha$	2.4168 $\pm$ 0.0553	1.5714 $\pm$ 0.0328	1016 sec
Regional adaptive fusion	$\alpha_{i}$	2.6208 $\pm$ 0.0672	1.7414 $\pm$ 0.0442	1140 Sec
Depth-dependent fusion	$\alpha_{\ell}$	2.4528 $\pm$ 0.0570	1.5883 $\pm$ 0.0307	1068 Sec
Joint regional-depth fusion	$\alpha_{i,\ell}$	2.6324 $\pm$ 0.0782	1.7373 $\pm$ 0.0466	1090 Sec

Table 7: Additional comparison of different fusion granularities under the same experimental setting on the SRED Greenland benchmark. The globally shared

\alpha

is the fusion strategy used in the final K-STEMIT model.

Run	K-STEMIT		KI Multi-Branch		GAT + Physical Features		PI-GraphSAGE-LSTM		ST-GRIT
Run	RMSE	MAE	RMSE	MAE	RMSE	MAE	RMSE	MAE	RMSE	MAE
1	2.4028	1.5448	2.4430	1.5899	2.5764	1.7099	2.7972	1.9495	2.8252	1.9167
2	2.4168	1.5958	2.4924	1.6399	2.5732	1.7177	2.9519	2.0500	2.9153	1.9253
3	2.5214	1.6237	2.5700	1.6619	2.7110	1.7663	3.0394	2.1365	2.9683	1.9468
4	2.3671	1.5399	2.4151	1.5808	2.5289	1.6775	2.8534	1.9766	2.8190	1.8568
5	2.3760	1.5530	2.5254	1.6148	2.5432	1.6940	2.9095	2.0581	2.9054	1.9250

Table 8: Per-run RMSE and MAE values used in the paired statistical tests. KI denotes knowledge-informed and PI denotes physics-informed.

Comparison	Metric	Mean Diff.	$t$	$p$ ( $t$ -test)	$W$	$p$ (Wilcoxon)
K-STEMIT vs KI Multi-Branch	RMSE	0.0724	3.5872	0.0230	0.0	0.0625
K-STEMIT vs KI Multi-Branch	MAE	0.0460	11.1466	$3.69\times 10^{-4}$	0.0	0.0625
K-STEMIT vs GAT + physical features	RMSE	0.1697	29.6171	$7.74\times 10^{-6}$	0.0	0.0625
K-STEMIT vs GAT + physical features	MAE	0.1416	20.4686	$3.36\times 10^{-5}$	0.0	0.0625
K-STEMIT vs PI-GraphSAGE-LSTM	RMSE	0.4935	18.7834	$4.73\times 10^{-5}$	0.0	0.0625
K-STEMIT vs PI-GraphSAGE-LSTM	MAE	0.4627	22.5506	$2.29\times 10^{-5}$	0.0	0.0625

Table 9: Paired statistical comparison of K-STEMIT against representative baselines over five matched runs. Positive mean difference indicates lower error for K-STEMIT, computed as baseline error minus K-STEMIT error. Because the paired

t

-tests were conducted in the order (baseline, K-STEMIT), a positive

t

-statistic also indicates lower error for K-STEMIT.

Figure 6: Training and validation loss curves of K-STEMIT over all five data splits. Across splits, the losses show overall downward trends with only minor transient fluctuations, indicating stable convergence behavior during optimization.

Figure 7: Violin plots of per-image RMSE distributions across five runs for K-STEMIT, Multi-Branch Graph Neural Network, and Physics-Informed GraphSAGE-LSTM. The figure provides a distributional view of prediction errors beyond the averaged RMSE values reported in the main text.

Figure 8: Taylor diagram summarizing the agreement between model predictions and ground truth across five runs for K-STEMIT, Multi-Branch Graph Neural Network, and Physics-Informed GraphSAGE-LSTM in terms of correlation, standard deviation, and centered RMSE. Points closer to the reference point indicate better overall agreement.

Figure 9: Comparison between qualitative results of different graph models. The blue line is used to generate the graphs. The green line is the groundtruth (manually-labeled ice layers) and the red line is the model prediction.

Figure 10: Relative MAE on individual ice layers for each model.

Figure 11: Centered cosine-similarity heatmaps of node representations from the raw input, GraphSAGE spatial branch, temporal convolution branch, and fused representation for a representative held-out sample. Each row/column corresponds to one radargram column node.

Figure 12: PCA visualization of node representations colored by radargram column index for the same held-out sample. Each point denotes one node, which is one fixed radargram column.

Figure 13: PCA visualization of node representations colored by amount of snow pack for the same held-out sample.

Figure 14: ( a ) Exact Shapley values computed from the complete 32 ablation table and averaged across the five test splits. Larger positive values indicate greater average reduction in prediction error relative to the empty-feature baseline. ( b ) Mean test-set feature-feature Spearman correlation matrix among the five MAR variables, showing redundancy in the raw physical inputs. ( c ) Mean pairwise interaction scores derived from the same coalition-value table used for the Shapley analysis; positive values indicate synergistic interactions and negative values indicate sub-additive or redundant interactions. Taken together, these panels distinguish marginal association, input redundancy, and conditional predictive contribution.

Feature	Test pooled $\rho$	Test within-radargram $\rho$
Snow Mass balance (SMB)	$0.9802\pm 0.0019$	$0.2492\pm 0.0379$
Average Yearly Surface Temperature(Surface Temp)	$0.8963\pm 0.0031$	$0.1408\pm 0.0185$
Height Change Due To Refreezing (Refreezing)	$0.3915\pm 0.0382$	$-0.0354\pm 0.0335$
Height Change Due To Melting (Melt)	$0.4191\pm 0.0399$	$-0.0418\pm 0.0359$
Snowpack Height (Snowpack)	$0.0337\pm 0.0558$	$-0.0001\pm 0.0219$

Table 10: Summary of feature-thickness correlation analysis on the five test splits. “Pooled” Spearman correlation is computed over all aligned

(\text{radargram},\text{node})

pairs after averaging each MAR feature across the five observed upper layers and summarizing the 15-layer target into one node-level thickness scalar. “Within-radargram” Spearman reports the mean correlation computed separately within each radargram across its 256 nodes.

Number of Predicted Layers	All Layer RMSE	All Layer MAE
5	1.7716 $\pm$ 0.1020	1.1634 $\pm$ 0.0747
8	1.8359 $\pm$ 0.1048	1.2274 $\pm$ 0.0689
10	1.7973 $\pm$ 0.1191	1.2240 $\pm$ 0.0828
13	1.8103 $\pm$ 0.1160	1.2504 $\pm$ 0.0802
15	1.8724 $\pm$ 0.0949	1.2982 $\pm$ 0.0609
18	1.8838 $\pm$ 0.1251	1.3143 $\pm$ 0.0843

Table 11: Generalization of K-STEMIT to different number of prediction layers with a fixed input setting of

m=5

observed layers. All experiments use the same subset of radargrams containing at least

5+21

complete layers, so that changing the number of predicted layers

n

does not change the training/test sample population. For smaller

n

, only the uppermost

n

target layers below the five observed input layers are used. Errors are reported as mean

\pm

standard deviation over repeated splits/runs.

Figure 15: Data availability and prediction-horizon generalization for K-STEMIT. ( a ) Number of radargrams containing at least a given number of complete layers. Because the number of available radargrams decreases sharply for very large complete-layer counts, we set the maximum prediction horizon to $n=21$ with $m=5$ observed input layers, which requires at least $26$ complete layers per radargram. This choice balances deeper-layer evaluation with sufficient training data. ( b ) Relative MAE for different target-layer years under prediction horizons $n\in\{5,8,10,13,15,18,21\}$ . Older years correspond to deeper internal layers, which are generally more challenging because radar signal quality and layer interpretability decrease with depth.

Figure 16: Qualitative examples under two types of incomplete deeper target-layer annotation: (a) some deeper target layers are partially missing, and (b) some deeper target layers are entirely missing. In both cases, the five upper observed layers are used as inputs to K-STEMIT, which predicts the deeper layers across the radargram. Where target layer annotations are available, the predictions can be directly compared with the annotated layers; where annotations are missing, the model still produces spatially continuous layer estimates.

Spatiotemporal GNNs

Spatiotemporal graph neural networks model how signals and relationships evolve across connected locations, agents, or physical systems.

See all 147 papers

Physics-Conditioned Synthesis of Internal Ice-Layer Thickness for Incomplete Layer Traces
Apr 22, 2026Zesheng Liu, Maryam RahnemoonfarIncomplete Data Imputation Remote Sensing
Department of Computer Science and Engineering Lehigh University Bethlehem, USA
ICIBS: Increasing the Context for Ice Boundary Segmentation
Oct 20, 2026Marcel Dreier, Dakota Pyles, Nora Gourmelon +6Remote Sensing Image Segmentation Remote Sensing
Pattern Recognition Lab, Computer Science Department, Friedrich-Alexander-Universität Erlangen-Nürnberg, Erlangen, Germany · Institut für Geographie, Department of Geography and Geosciences, Friedrich-Alexander-Universität Erlangen-Nürnberg, Erlangen, Germany
Neural Posterior Estimation of Terrain Parameters from Radar Sounder Data
May 5, 2026Jordy Dal Corso, Annalena Kofler, Marco Cortellazzi +2Simulation-Based Inference Parameter Estimation
University of Trento, Trento, Italy · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Max Planck Institute for Gravitational Physics (Albert Einstein Institute), Potsdam, Germany +1

Apr 22, 2026J/K move · Enter open · S save

Physics-Conditioned Synthesis of Internal Ice-Layer Thickness for Incomplete Layer Traces

Zesheng Liu, Maryam Rahnemoonfar

Department of Computer Science and Engineering Lehigh University Bethlehem, USA

Incomplete Data Imputation Remote Sensing

Open paper AskPDF arXiv

Internal ice layers imaged by radar provide key evidence of snow accumulation and ice dynamics, but radar-derived layer boundary observations are often incomplete, with discontinuous traces and sometimes entirely missing layers, due to limited resolution, sensor noise, and signal loss. Existing graph-based models for ice stratigraphy generally assume sufficiently complete layer profiles and focus on predicting deeper-layer thickness from reliably traced shallow layers. In this work, we address the layer-completion problem itself by synthesizing complete ice-layer thickness annotations from incomplete radar-derived layer traces by conditioning on colocated physical features synchronized from physical climate models. The proposed network combines geometric learning to aggregate within-layer spatial context with a transformer-based temporal module that propagates information across layers to encourage coherent stratigraphy and consistent thickness evolution. To learn from incomplete supervision, we optimize a mask-aware robust regression objective that evaluates errors only at observed thickness values and normalizes by the number of valid entries, enabling stable training under varying sparsity without imputation and steering completions toward physically plausible values. The model preserves observed thickness where available and infers only missing regions, recovering fragmented segments and even fully absent layers while remaining consistent with measured traces. As an additional benefit, the synthesized thickness stacks provide effective pretraining supervision for a downstream deep-layer predictor, improving fine-tuned accuracy over training from scratch on the same fully traced data.

Abstract

Spatiotemporal GNNs

Explore similar work

Physics-Conditioned Synthesis of Internal Ice-Layer Thickness for Incomplete Layer Traces

ICIBS: Increasing the Context for Ice Boundary Segmentation

Neural Posterior Estimation of Terrain Parameters from Radar Sounder Data