goia-impute-base

Zero-shot probabilistic imputation for univariate time series. A 10.2M-parameter bidirectional encoder (ModernBERT-style, one point per position) that fills gaps in any numeric series without per-dataset training, returning 9 quantiles per imputed position.

This is an experiment in general-purpose neural imputation rather than a production system. In short: it consistently beats classical interpolation on block gaps and mixed missingness patterns, produces well-calibrated intervals without post-hoc calibration, and adds little over interpolation on scattered point gaps in high-frequency data.

Usage

The weights ship as a self-contained safetensors file (model config embedded in the metadata). Inference code lives in the project repository and depends only on torch, numpy and safetensors:

import numpy as np
from inference import Imputer   # from the project repo

imp = Imputer.from_safetensors("goia-impute-base.safetensors")

s = np.asarray(my_series, dtype=np.float32)   # NaNs mark the gaps
res = imp.impute(s)

res.filled        # series with NaNs replaced by the predictive median
res.median        # median at imputed positions
res.q(0.1)        # any quantile in (0.1 ... 0.9)
res.missing_idx   # imputed positions

Series longer than the 2048-point context are handled with overlapping sliding windows. An extra_mask argument lets you hide observed positions (evaluation, what-if analysis). imp.impute_df(series) gives a pandas DataFrame with filled, q10, q50, q90.

Code repository: https://github.com/gredio/goia-impute

Model

Bidirectional encoder, no tokenizer, no patching: each timestep is one sequence position with two input channels (normalized value, visibility flag) plus a learned additive [MASK] embedding at hidden positions — no sentinel values. ModernBERT-style body: pre-norm, RoPE, GeGLU MLPs, no biases, alternating attention (1 global layer every 3, the rest block-local with window 128; all calls take the Flash Attention path). RevIN statistics are computed only over observed-and-visible positions, so normalization never sees the values to be predicted. A 9-quantile head is trained with pinball loss on hidden-and-originally-observed positions only.

12 layers, d=256, 8 heads, GeGLU hidden 768, max context 2048. 10.23M parameters.

Training

MLM-style with a masking curriculum: one regime per window (scattered points 35%, contiguous blocks 35%, mixed 15%, tail 10%, head 5%), BERT-style 80/10/10 split over hidden positions. Native NaNs in the corpus are treated as real missingness, never pre-interpolated.

Data: ~3.87M series — 90% processed GiftEvalPretrain, 7% TempoPFN, 3% KernelSynth; only the train portion of each series. Two phases: 250k steps at window 1024 (batch 128, lr 3e-4, cosine), then 25k context-extension steps at 2048 (batch 64, constant lr 3e-5). Single RTX 4090, bf16.

Evaluation

Reproducible protocol over all 97 GIFT-Eval configurations: deterministic masks generated from a published seed (EVAL_SEED = 2026), cycling five missingness regimes, scoring only hidden-and-originally-observed positions. Ratios are per-series against linear interpolation, aggregated with a geometric mean (< 1 means the model wins).

Among the classical baselines evaluated, the model achieves the lowest MAE on 77% of configurations; this evaluation does not cover supervised neural imputation methods. The bootstrap CI was computed for the primary metric (relative MAE). Per missingness regime (relative MAE): points 0.82, blocks 0.69, mixed 0.68, tail 0.73, head 0.80 — the pattern is that where immediate neighbors are highly informative (point gaps, minute-level frequencies), interpolation is a very strong baseline; on long or structured gaps the model exploits global series structure.

Forecasting as a special case. Tail imputation is forecasting, so the model was also run through the official GIFT-Eval forecasting benchmark: normalized CRPS 0.576 (geomean over 97 configs, seasonal naive = 1), on par with chronos_bolt_small (0.577) and behind size-comparable forecasting specialists such as FlowState-9.1M (0.502). It is competitive at short horizons (0.557) and degrades at medium/long horizons where autoregressive extension kicks in. If forecasting is your goal, dedicated models of the same size do better.

Reproducibility

Evaluation masks are generated by a deterministic RNG from EVAL_SEED = 2026, cycling five missingness regimes. The generator state advances series by series, so exact reproduction requires the published masking code, the seed, and the same list of series in the same order (with a pinned numpy version). Under those conditions, third parties can regenerate the exact masks and evaluate any other model on identical targets.

Limitations

• Univariate: cross-channel correlations are not used.
• Small margin over pchip/linear on scattered point gaps in high-frequency data.
• Medium/long-horizon forecasting trails specialized models.
• Windows with fewer than ~8 visible points cannot be reliably normalized and are not scored.
• Domains far from the training mix may behave worse than the benchmark numbers suggest.

Intended use

Filling gaps in numeric time series (sensors, energy, weather, traffic, business metrics) where a probabilistic estimate with calibrated intervals is preferable to point interpolation. Not intended for high-stakes decisions without domain-specific validation.