Filling the G_ap_s: Multivariate Time Series Imputation by Graph Neural Networks

In plain words

Sensors that feed a time series often drop out, so some readings are simply missing. The sensors are also related to each other, as nearby stations or road sensors often are. This paper asks whether a model can fill the gaps by looking at both the past and the future of a sensor and at what its neighbours are reporting at the same moment.

The answer is a network that treats the sensors as nodes in a graph. At each time step it makes a rough guess for each missing value, lets every sensor pass information to its neighbours, and then uses that neighbourhood information to refine the guess. It runs this once forward in time and once backward, and a small final network combines the two.

The authors report that this beats earlier deep-learning imputers on several real datasets, mostly on air quality and traffic. The results are the authors' own.

TL;DR

• Imputes missing values in a multivariate series whose channels sit on a fixed graph (sensors plus adjacency), using a bidirectional recurrent network whose gates are graph message-passing layers §4.
• Each direction has a spatio-temporal encoder and a spatial decoder that imputes twice per step: once from the hidden state, once from neighbours §4.
• Author-reported: lowest MAE in every real-data table row. Largest reductions are in the point-missing traffic and smart-meter settings (about 54% to 55% below BRITS) tab. 2; out-of-sample AQI is about 27% below tab. 1.
• Relevant if your channels have a known or computable similarity graph and the graph is fixed over time. Dense N×N supports in the released code scale poorly with node count spatial_conv.py.

Problem

Missing values are common in sensor networks because sensors fail independently, so at any step only some channels are usually affected §1. Imputation should use temporal dependencies and the relations between sensors, but the deep-learning imputers the authors review either ignore relations or bolt simple changes onto sequence models §1 §2. The paper states that no earlier GNN method targeted missing-value imputation for generic multivariate series §2.

The formal setting is a sequence of graphs with a binary mask marking observed entries. The method assumes stationary missingness, mostly under MAR, and makes no assumption on the number or length of missing blocks §3. The paper distinguishes two settings. In-sample imputation trains and evaluates on the same sequence. Out-of-sample imputation trains and evaluates on disjoint sequences §3.

Method

GRIL has two units of computation, one per time direction. Each unit runs a spatio-temporal encoder (a GRU whose gates are message-passing networks) and a spatial decoder (two rounds of imputation). A final MLP combines the forward and backward representations §4.

Message passing. A generic message-passing layer updates each node from its own features and an aggregate over neighbours. The aggregate is permutation-invariant, with sum or mean as examples §4, eq. 2. The paper uses diffusion convolution [atwood2016diffusion] as the message operator for all experiments §5 app. A.

Recurrent gates §4, eq. 3-6. The GRU reset gate r, update gate u and candidate state c are each computed by a message-passing network over the concatenation of the previous second-stage imputation at the node, its mask, and its hidden state. The reset gate multiplies the previous hidden state before the candidate is formed. The new state is h = u * h_prev + (1 - u) * c. The gate networks have separate weights; the equations write them alike.

Filler operator §4, eq. 8. Phi(Y) = M * X + (1 - M) * Y. Observed values are kept, and the model's prediction fills the gaps.

First stage §4, eq. 7-8. A linear readout of the previous hidden state gives a prediction for every node. The filler replaces only the missing entries of the current input with that prediction.

Spatial decoder, stage two §4, eq. 9-10. For each node, a message-passing step aggregates from its one-hop neighbours only (not itself) the filled first-stage values, their masks and their hidden states. The result is concatenated with the node's previous hidden state, passed through a second linear readout, and filled again. Because the node's own input is excluded, the representation does not depend on that node's own input, which the paper argues regularises the model toward spatial information §4.

Recurrence. The second-stage imputed input updates the hidden state, and the process moves to the next step §4.

Bidirectional output §4, eq. 11. The forward and backward units are run on the sequence and its time reverse. For each node and step, an MLP takes the forward spatial representation and hidden state and the backward ones, and outputs the final value. The output is filled as in eq. 8 §4.

Training objective app. A. The loss sums MAE over five terms: the final output, plus the first-stage and second-stage outputs of each direction. It is computed on the originally observed entries, not on the entries hidden from the input.

Pseudocode of one unidirectional pass, following the paper (not the code):

# One GRIL direction over T steps  [§4, eq. 3-10]
# Input:  X (B, T, N, d), observed mask M (B, T, N, d), adjacency W (N, N)
# Output: per step, first-stage X1, second-stage X2, representation R
def gril_direction(X, M, W):
    # Hidden state: learned per-node init or a constant  [§4]
    H = init_hidden_state(B, N, hidden)     # (B, N, hidden)
    for t in range(T):
        H_prev = H
        # First stage: linear readout of the hidden state  [§4, eq. 7]
        Y1 = linear_readout_1(H_prev)       # (B, N, d)
        # Keep observed values, fill the gaps with the prediction  [§4, eq. 8]
        X1 = fill(Y1, X[t], M[t])           # (B, N, d)
        # Spatial decoder: one-hop neighbour messages, no self  [§4, eq. 9]
        S = neighbour_messages(X1, M[t], H_prev, W)  # (B, N, hidden)
        # Second linear readout over [S || H_prev], then fill  [§4, eq. 10]
        Y2 = linear_readout_2(concat(S, H_prev))  # (B, N, d)
        X2 = fill(Y2, X[t], M[t])           # (B, N, d)
        # Keep what the bidirectional head needs  [§4, eq. 11]
        R[t] = concat(S, H_prev)            # (B, N, 2 * hidden)
        # Update the recurrent state from the completed input  [§4, eq. 3-6]
        H = gated_update(X2, M[t], H_prev, W)  # (B, N, hidden)
    return X1, X2, R

Bidirectional head and objective (the paper writes the head as in eq. 11):

# Bidirectional head and training loss  [§4, eq. 11] [app. A]
# Input:  X, M (B, T, N, d); W (N, N)
def gril_impute(X, M, W):
    # Forward pass over time, then backward pass on the reversed sequence
    fwd = gril_direction(X, M, W)
    bwd = gril_direction(reverse_time(X), reverse_time(M), W)
    # Per node and step: MLP over both directions' representations  [eq. 11]
    Y = mlp(concat(fwd.R, reverse_time(bwd.R)))  # (B, T, N, d)
    return fill(Y, X, M)                    # observed values kept  [§4]

def training_loss(X, M, W):
    # Randomly hide an extra 5% of observed inputs per batch  [app. A]
    M_in = bernoulli(M * 0.95)              # (B, T, N, d)
    # Sum MAE over the five outputs, on originally observed entries  [app. A]
    return sum_of_mae(all_five_outputs(X, M_in, W), X, M)

The code adds the five terms with weight 1 graphfiller.py.

Key equations, with the numbers the paper's LaTeX source assigns (consecutive through the paper; checked against the source, not a compiled PDF):

\mathrm{MPNN}_k(z^i, W) = \gamma_k\Big(z^i, \sum_{j \in N(i)} \rho_k(z^i, z^j)\Big) eq. 2

\hat{Y}^{(1)}_t = H_{t-1} V_h + b_h eq. 7

\Phi(Y_t) = M_t \odot X_t + \overline{M}_t \odot Y_t eq. 8

\hat{X}^{(2)}_t = \Phi\big( [S_t \,\|\, H_{t-1}] V_s + b_s \big) eq. 10

\hat{y}^i_t = \mathrm{MLP}\big( [s^{i,\mathrm{fwd}}_t \,\|\, h^{i,\mathrm{fwd}}_{t-1} \,\|\, s^{i,\mathrm{bwd}}_t \,\|\, h^{i,\mathrm{bwd}}_{t+1}] \big) eq. 11

w^{i,j} = \exp\!\left(-\frac{\mathrm{dist}(i,j)^2}{\gamma}\right) \text{ if } \mathrm{dist}(i,j) \le \delta, \text{ else } 0 eq. 12, app. B.1

The released code is the implementation of these equations [code: lib/nn/layers/gril.py#L11-L54] gril.py. Differences are in Paper vs code.

Hyperparameters

Values are from the paper unless the source column says code. Hyperparameters are shared across datasets app. A.

Results

Author-reported. Each number is a mean over 5 runs unless noted. Hardware is not stated. tab. 1 reports the air-quality datasets; tab. 2 reports traffic and smart grids. Baseline column shows BRITS, the paper's main competitor §5. The percentage column is computed here from the table values, not reported by the paper.

Other reported results:

• Ablation (AQI / METR-LA block / CER-E point, MAE): GRIL 14.73 / 2.03 / 0.29; without spatial decoder 15.40 / 2.32 / 0.29; with denoising decoder 17.23 / 2.96 / 0.32; MPGRU 18.76 / 2.57 / 0.41 tab. 3 tab. 8.
• Graph structure (METR-LA block, MAE): learned graph 2.03; fully connected 2.63; no edges 3.42; BRITS 2.34 tab. 9.
• Matrix factorization with side information (AQI-36, MAE): TRMF 15.97, GRMF 19.29, MF 30.54 tab. 7. GRIL in-sample is 10.51 tab. 7.
• Sensitivity (METR-LA point, MAE at 10% to 90% masked): GRIL 1.87 to 2.84, BRITS 2.32 to 4.02 tab. 10.

Paper vs code

Repository: the paper's footnote links https://github.com/Graph-Machine-Learning-Group/grin at commit 4a28afbb092600b6e6abeeabaaf67e87dbd1ed6e app. B. The README lists the authors. The repository has no LICENSE file; the code is not quoted here and is only cited by path and line.

• Name. The paper calls the model GRIL. The code and README use GRIN (lib/nn/models/grin.py) grin.py.
• Kernel width. The paper sets gamma to the standard deviation of AQI-36 distances in eq. 12 app. B.1. The code computes exp(-(dist / theta)^2) with theta = std utils.py, so the effective denominator is std squared air_quality.py.
• Adjacency cutoff. The paper cuts edges beyond about 40 km for AQI app. B.1. The code drops edges whose kernel weight is below 0.1 utils.py run_imputation.py. No distance cutoff is coded. For the traffic sets the code uses sigma = std(finite distances) and the same 0.1 weight threshold pems_bay.py metr_la.py.
• Validation split. The paper states 70 / 10 / 20 for traffic and CER-E app. B.2 app. B.3. The code takes 20% for test and then 10% of the remaining 80% for validation, which is 8% of the total, so 72 / 8 / 20 pems_bay.py metr_la.py with the default --val-len 0.1 run_imputation.py.
• Node embeddings. The paper does not mention learned node embeddings for the main model. The configs for real data set d_emb = 8 air36.yaml, and the imputation MLP concatenates the embedding and the mask to the two representations gril.py. Eq. 11 omits both. Synthetic data uses d_emb = 0 synthetic.yaml.
• Decoder layers. The paper describes one message-passing step, then a linear readout §4, eq. 9-10. The code wraps the graph convolution in a 1x1 input projection, a PReLU, and a 1x1 output projection before the readout gril.py gril.py gril.py.
• Aggregation. The paper allows sum or mean §4. The code uses row-normalised adjacency (each row divided by its sum, with epsilon 1e-8), so messages are weighted means spatial_conv.py [code: lib/init.py#L14].
• Encoder diffusion. The encoder uses both the forward and transposed backward supports, with powers up to order 2, and includes the node itself spatial_conv.py gcrnn.py. The decoder excludes self-loops gril.py.
• Optimizer and schedule. The paper gives a cosine schedule from 0.001 over 300 epochs and does not name the optimizer app. A. The code uses Adam with cosine annealing down to eta_min = 1e-4 run_imputation.py.
• Gradient clipping. Not stated in the paper. The code clips the gradient norm to 5 run_imputation.py.
• Air out-of-sample. For the air data, values held out for evaluation are added back to the observed mask over the training slice, so training sees them run_imputation.py. The paper describes out-of-sample training on disjoint sequences without this detail §3.
• Evaluation output. At test time, observed entries are copied from the input, so the model only fills gaps grin.py. The paper does not say this.
• CER-E loader. The configs irish_block and irish_point exist irish_block.yaml, and run_baselines.py lists irish run_baselines.py. run_imputation.get_dataset has no branch for CER-E run_imputation.py, and no loader is in lib/datasets. The README says CER-E must be obtained separately. The CER-E results therefore cannot be rerun from this repository alone as released.
• Dataset and metric names. The paper's "CER-E" is "Irish" in the code; the paper's "AQI-36" is air36 run_imputation.py.

No discrepancies were found for: the MPNN gates and filler form, the five-term loss (weight 1) graphfiller.py, the 64-unit hidden size, the one-step decoder order, window lengths for AQI, AQI-36 and traffic, the 5% extra masking filler.py, and the early-stopping patience of 40 epochs.

Pitfalls

• Which equation to implement. Eq. 11 has no mask and no embedding; the code has both. The decoder has extra projections the paper does not write down. Follow the code if you want the reported numbers, and say so in your notes.
• Kernel and threshold. Read gamma and delta as the code does, or your graph will differ from the paper's. The paper's delta (about 40 km) does not appear in the code; the 0.1 weight cutoff does.
• Table 6 and the text disagree. The text says BRITS is "more accurate in the setting with varying charge". Table 6 shows GRIL ahead in both settings (0.0530 vs 0.1089 MAE with varying charge). The sentence probably means BRITS does better with varying than with fixed charge. Trust the table.
• Abstract claim. "Often higher than 20%" holds for AQI out-of-sample (27%) and for several point-missing settings, but not everywhere: METR-LA is 13% to 18% lower, and AQI-36 out-of-sample is 17%. Do not treat 20% as a guarantee.
• Two evaluation protocols. Out-of-sample error in tab. 1 is averaged over windows (per-batch test metrics). The code also reports errors on stitched series with mean aggregation run_imputation.py. The two numbers differ, so use the one your comparison needs.
• Training loss includes dropped inputs. The loss compares against the originally observed mask, so entries that were randomly hidden from the input are still reconstructed in training graphfiller.py graphfiller.py.
• Hidden-state init. The paper allows a constant or a learnable embedding §4. The code learns a per-node initial state drawn from a normal with standard deviation 1/sqrt(d_hidden) gril.py.
• Dense graph cost. The released code uses dense N×N supports and einsum over all nodes at every step spatial_conv.py. The paper says message passing costs O(E) with a sparse implementation or O(N²) dense, and that the overhead can be amortised over parallel computation app. C.3. The released code does not use a sparse path.
• Fixed topology. The method assumes W_t = W for all t §3. Dynamic graphs need an adaptation the paper does not give.
• Stationarity. The paper assumes stationary missingness and mostly MAR §3. Non-stationary missingness is listed as future work §6.
• Scale of results. Five runs, author-reported, on a small set of benchmarks. Baselines were tuned by the authors; BRITS hidden sizes were tuned on validation app. A. Hardware and total training time are not stated.
• Virtual sensing. The numbers in §5.2 are reported without a baseline. Treat them as a qualitative demonstration.

Lineage

• Builds on: gated recurrent units (Cho et al., 2014) with message-passing gates §4, eq. 2-6; diffusion convolution (Atwood and Towsley, 2016) as the message operator §5. No cards exist for these yet.
• Related: the MPGRU baseline is described as similar to DCRNN (Li et al., 2018) §5. Graph-based spatio-temporal forecasting is discussed in §2.
• Compared against: BRITS (Cao et al., 2018), the principal baseline, plus rGAIN, VAR, MICE, MF, KNN, mean imputation and MPGRU §5 tab. 1 tab. 2.
• Superseded by: not known from this wiki.