google-research / google-research/flood-forecasting
[Bug] HandoffForecastLSTM: Silent covariate shift caused by dictionary iteration order in feature concatenation
- 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
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