Technical Summary: StruMPL

Problem Formulation

The paper addresses the challenge of estimating forest Aboveground Biomass (AGB) and associated structural variables (canopy height, cover, stem density, wood density) from Earth Observation (EO) data. The problem is characterized by three coupled, structurally incompatible constraints that standard multi-task learning (MTL) and semi-supervised methods fail to address jointly:

1. Heterogeneous Disjoint Partial Supervision: Training labels come from two distinct sources with non-overlapping variable sets. Spaceborne lidar (e.g., GEDI) provides structural metrics (height, cover) at millions of locations but no biomass. Ground-based plots provide biomass and some structural metrics at thousands of locations but no canopy structure at the lidar scale. No single training sample contains labels for all target variables.
2. Missing Not At Random (MNAR) Labels: Ground plot locations are not a random sample of the forest population; surveyors avoid inaccessible terrain and degraded stands. Consequently, observed AGB is biased toward intermediate values, leading standard models to systematically underestimate biomass in dense forests.
3. Inter-Task Physical Constraints: Target variables are linked by known allometric laws (e.g., $AGB \approx g(H, C, SD, WD; \phi)$). However, because the training data never observes all variables simultaneously (due to the disjoint sources), these parameters $\phi$ cannot be fitted directly from the data alone.

The authors formalize this as a multi-task dense regression problem under MNAR labels with inter-task physical constraints, proposing StruMPL (Structured Multi-task Physics-constrained Learning) as a unified solution.

Methodology

StruMPL employs a single-objective framework with a shared encoder and specialized heads to address the three problem components simultaneously.

Architecture

• Shared Encoder: A ResUNet with channel and spatial attention processes multi-channel EO inputs (e.g., SAR, optical, coordinates) to produce a feature map $z$.
• Regression Heads: $K$ independent convolutional heads predict each target variable $\hat{y}_k$ in z-score space.
• Imputation Heads: $K$ independent heads predict an imputation baseline $\hat{m}_k$ (the expected outcome given covariates) for every pixel.
• Propensity Head: A head estimates the per-pixel, per-variable observation probability $\hat{\pi}_k$ (the propensity score).
• Physics Module: A learnable, differentiable allometric equation $g_\phi$ that evaluates the physical relationship between structural predictions and biomass predictions at every pixel.

Loss Function

The total objective $L_{total}$ combines five terms:
$$L_{total} = L_{sup} + \lambda_{phys} L_{phys} + \lambda_{cons} L_{cons} + \lambda_{bias} L_{bias} + \lambda_{imp} L_{imp}$$

1. Supervised Loss ($L_{sup}$) with AIPW:
   Instead of standard masked MSE, the paper uses the Augmented Inverse Probability Weighting (AIPW) estimator. The pseudo-outcome $\tilde{y}$ is defined as:
   $$\tilde{y} = \mu + R \frac{1}{\hat{\pi}}(y - \mu)$$
   where $\mu$ is the imputation estimate and $R$ is the observation mask.

   • Critical Design: Two stop-gradients are applied: one on the propensity $\hat{\pi}$ and one on the imputation $\mu$.
     • Detaching $\hat{\pi}$ prevents the encoder from inflating propensities to 1.0 to minimize loss (propensity collapse).
     • Detaching $\mu$ prevents the imputation heads from collapsing onto the regression heads, preserving the variance-reduction property of AIPW over standard IPW.
   • This loss corrects for MNAR bias by re-weighting observed residuals while using the imputation model to provide supervision at unlabelled pixels.

2. Physics Consistency Loss ($L_{phys}$):
   This term enforces the allometric relationship $g_\phi$ on the model's own predictions at every pixel (labelled or unlabelled). It acts as a semi-supervised signal, forcing structural predictions (from GEDI-supervised pixels) to be consistent with biomass predictions (from plot-supervised pixels) via the learned allometric coefficients. The parameters $\phi$ are learned jointly with the network.

3. Auxiliary Losses:

   • $L_{bias}$: Binary cross-entropy training the propensity head against the observation mask.
   • $L_{cons}$: Consistency loss under input augmentation (FixMatch-style) to regularize the encoder.
   • $L_{imp}$: Masked MSE training the imputation heads on labelled pixels only.

Training Strategy

The authors employ source-balanced batching, constructing each batch as a 50/50 mix of GEDI and plot samples. This ensures the encoder receives balanced gradients from both label structures and allows the physics loss to operate on batches containing both structural and biomass labels.

Key Contributions

1. Problem Formalization: The paper defines a specific class of multi-task dense regression problems involving heterogeneous disjoint partial supervision, MNAR labels, and inter-task physical constraints.
2. Joint Propensity & AIPW Design: It introduces a joint propensity model for spatial MNAR correction in dense regression. Crucially, it demonstrates analytically and empirically that two stop-gradients (on propensity and imputation) are individually necessary for stable joint training. Without them, the framework collapses to biased MSE or unstable IPW.
3. Learnable Physics Module: A differentiable constraint module evaluates known biophysical relationships on model predictions at every pixel. Unlike standard Physics-Informed Neural Networks (PINNs) with fixed parameters, StruMPL learns the allometric coefficients jointly, allowing adaptation to specific biomes and satellite-scale scaling.
4. Empirical Validation: The framework is validated on two ecologically distinct biomes (Mediterranean forests in Spain and dry tropical/savanna forests in Africa) using identical architecture and training recipes, demonstrating portability across variable availability and source imbalance.

Results

The framework was evaluated on the AGB target variable (Mg/ha) across two sites:

• Performance: Full StruMPL achieved the lowest RMSE and bias on both sites compared to single-output baselines, ablation variants, and the closest published method.
  • Spain (SNFI): RMSE 41.9 ± 0.5 Mg/ha (vs. 43.1 for the strongest single-output baseline).
  • Africa (SEOSAW): RMSE 20.7 ± 0.6 Mg/ha (vs. 22.5 for single-output baseline).
• Ablation Insights:
  • MNAR Correction: Stratified analysis on the Spain dataset showed that the AIPW correction specifically reduced bias in the high-AGB tail (the under-represented regime) by approximately 54% compared to naive MSE.
  • Necessity of Stop-Gradients: Removing either stop-gradient led to "propensity collapse" (mean propensity $\to$ 1.0) or "imputation collapse" (imputation $\to$ regression), degrading performance significantly.
  • Physics Module: The ecologically motivated allometric form outperformed both a plain power law and a learned MLP, suggesting the inductive bias of the allometric form is more valuable than raw expressiveness.
• Learned Coefficients: The learned allometric exponents were small, indicating that at the satellite scale, biomass is largely explained by an additive baseline and a multiplicative scale, with structural variables acting as small corrections. The coefficients varied qualitatively between biomes (e.g., higher canopy-cover exponent in Spain vs. Africa), aligning with ecological intuition.

Significance and Claims

The paper claims that StruMPL successfully addresses a "perfect storm" of challenges in EO-based forest monitoring: disjoint data sources, severe class imbalance, and non-random sampling bias.

• Bias Reduction: The primary significance is the reduction of systematic bias in high-biomass regimes, which are critical for carbon stock estimation but typically under-represented in field data.
• Data Efficiency: By leveraging the physical constraints and the massive volume of structural data (GEDI) to supervise the biomass prediction, the framework effectively utilizes the sparse biomass labels.
• Generalizability: The authors suggest the formulation applies to any domain with structurally similar properties (e.g., hydrology, battery state estimation, multi-modal medical imaging) where target variables are linked by known laws but observed through disjoint, biased instruments.

The paper remains modest regarding theoretical guarantees, acknowledging that while AIPW offers double robustness under approximate ignorability, the shared encoder prevents formal cross-fitting guarantees. However, the empirical results confirm that the proposed architecture effectively reduces bias and improves accuracy in real-world, ecologically distinct settings.