google-research / google-research/flood-forecasting

[Bug] HandoffForecastLSTM: Silent covariate shift caused by dictionary iteration order in feature concatenation

Open
#281 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
343
Forks
87
PR merge metrics
No merged PRs in 30d

Description

### Description
In `googlehydrology.modelzoo.handoff_forecast_lstm.HandoffForecastLSTM`, feature concatenation relies on dictionary iteration order rather than the model's fixed feature list order. This creates a silent covariate shift risk where feature columns
can be permuted without raising PyTorch dimension errors.

### Locations
`googlehydrology/modelzoo/handoff_forecast_lstm.py`:
- `forward()` (lines 260–269)
- `save_state()` (lines 467–474)
- `predict_from_state()` (lines 490–497)

```python
hindcast_features = torch.cat(
[
t for f, t in data['x_d_hindcast'].items()
if f in self.hindcast_inputs
], dim=-1)
forecast_features = torch.cat(
[
t for f, t in data['x_d_forecast'].items()
if f in self.forecast_inputs
], dim=-1)
```
### Failure Mechanism

1. data['x_d_hindcast'].items() yields key-tensor pairs in the arbitrary insertion order of the input dictionary.
2. The model's weights in self.hindcast_embedding_net expect an immutable feature column order defined at initialization (e.g., column 0 = precipitation, column 1 = temperature).
3. If a different dataloader, dataset format (e.g. NetCDF vs Zarr), or pipeline creates data['x_d_hindcast'] with keys in a different order, columns swap positions (e.g., precipitation inputs get fed into temperature weights).
4. Because the tensor dimensions ([batch, seq_len, n_features]) remain identical, PyTorch raises no error, leading to silent model degradation.

### Reference Implementation

MeanEmbeddingForecastLSTM avoids this bug (googlehydrology/modelzoo/mean_embedding_forecast_lstm.py:553-556) by explicitly indexing the dictionary using the model's key sequence:

def _concat_tensors_from_dict(data: dict[str, torch.Tensor], *, keys: Iterable[str]) -> torch.Tensor:
return torch.cat([data[e] for e in keys], dim=-1)

### Proposed Fix

Iterate over self.hindcast_inputs and self.forecast_inputs instead of .items():

```python
hindcast_features = torch.cat(
[
data['x_d_hindcast'][f]
for f in self.hindcast_inputs
if f in data['x_d_hindcast']
],
dim=-1,
)

forecast_features = torch.cat(
[
data['x_d_forecast'][f]
for f in self.forecast_inputs
if f in data['x_d_forecast']
],
dim=-1,
)

Contributor guide

Open the contributing guide

Research direction

Start in googlehydrology/modelzoo/handoff_forecast_lstm.py and inspect forward(), save_state(), and predict_from_state(), then compare their concatenation logic with mean_embedding_forecast_lstm.py. Ensure each path uses the model’s hindcast_inputs and forecast_inputs order rather than dictionary iteration. Done means all three paths preserve fixed feature ordering when input dictionaries use different insertion orders.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
75/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.