Scaling Features Consistently Between Training and Inference

Stop training-serving skew in geospatial pipelines: fit scalers on training data only, persist the statistics as an artifact, and assert the contract at inference time.

Fit the scaler once, on the training split only. Save its statistics as a versioned artifact next to the model. At inference, load those statistics and apply them — never recompute. The rule sounds obvious and is broken constantly in geospatial pipelines, because inference runs tile by tile and computing a mean over the tile you already have in memory feels natural.

This page is about the contract between the two stages. For the choice of scaler — standard, robust, quantile — and the distributional reasons behind it, see Feature Scaling for Geospatial Inputs.

Re-fitting per tile makes a pixel’s value depend on its neighbours Two inference tiles are processed twice. With a re-fitted scaler, a forest tile and a mixed tile each have their own mean mapped to zero, so the same raw reflectance receives two different scaled values. With the training scaler loaded from an artifact, the same raw reflectance receives the same scaled value in both tiles. re-fit per tile load training stats forest tile raw NIR = 0.42 tile mean 0.42 scaled = 0.00 “perfectly average” mixed tile raw NIR = 0.42 tile mean 0.21 scaled = +1.90 “unusually bright” Same pixel, two answers. Predictions depend on batching. forest tile raw NIR = 0.42 train mean 0.24 scaled = +1.64 bright, as it should be mixed tile raw NIR = 0.42 train mean 0.24 scaled = +1.64 identical, as it should be Same pixel, one answer. Predictions are tile-independent.

Why This Fails in Geospatial ML Pipelines

Training-serving skew is the single most common cause of a geospatial model that scores well offline and disappoints in production, and scaling is where it usually enters. The mechanism is simple: a model does not learn “NIR reflectance above 0.4 means vegetation”, it learns “scaled feature 3 above 1.6 means vegetation”. If the scaling changes, the threshold now refers to a different physical quantity, and the model applies rules it learned about one distribution to another.

Geospatial inference makes re-fitting unusually tempting. Prediction runs tile by tile, and each tile arrives as a self-contained array. StandardScaler().fit_transform(tile) is one line and produces beautifully normalised inputs — and it means a pixel’s scaled value now depends on which pixels happened to share its tile. A tile of homogeneous forest has its mean mapped to zero, so every pixel in it reports as unremarkable; the same forest inside a mixed tile reports as strongly vegetated. Neither prediction is stable, and neither matches training.

The second route to skew is silent statistic drift in a shared preprocessing module. Someone updates the scaler to be fitted on “all available data” so a new region is covered; the training statistics change; the deployed model keeps its old learned thresholds. Nothing errors. The only signal is a slow degradation that looks exactly like the covariate shift described in model drift detection for geospatial inference — and is not drift at all, but a bug.

The third is nodata. If nodata sentinels reach the scaler during fitting, the mean and standard deviation are computed over a mixture of reflectance and -9999, so the statistics are dominated by pixels that are not data. Masking must happen before fitting, not after.

Core Principles

  • Fit on the training split only. Not on train plus validation, and never on the full dataset.
  • Persist the statistics as a versioned artifact. They are model state, exactly like weights.
  • Assert the feature order. A scaler is an array of numbers; a reordered column list applies band 3’s mean to band 5.
  • Mask nodata before fitting. Compute statistics over valid pixels only.
  • Log post-scaling moments at inference. Cheap, and it exposes skew immediately.
  • Fail closed on a missing artifact. A pipeline that silently falls back to re-fitting is worse than one that stops.

Production-Ready Code

from __future__ import annotations

import hashlib
import json
import logging
from dataclasses import dataclass

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class BandScaler:
    """Per-band standardisation whose statistics are an explicit, versioned artifact.

    Attributes:
        feature_names: Column order the statistics correspond to. Order is part of the contract.
        mean_, scale_: Per-feature statistics, fitted on the training split only.
    """
    feature_names: list[str]
    mean_: np.ndarray | None = None
    scale_: np.ndarray | None = None

    def fit(self, x: np.ndarray, nodata_mask: np.ndarray | None = None) -> "BandScaler":
        """Fit over valid pixels only. x has shape (n_samples, n_features)."""
        if x.shape[1] != len(self.feature_names):
            raise ValueError(
                f"expected {len(self.feature_names)} features, got {x.shape[1]}")
        valid = np.isfinite(x).all(axis=1)
        if nodata_mask is not None:
            valid &= ~nodata_mask
        if valid.sum() < 100:
            raise ValueError(f"only {int(valid.sum())} valid rows — refusing to fit")

        subset = x[valid]
        self.mean_ = subset.mean(axis=0).astype("float64")
        scale = subset.std(axis=0).astype("float64")
        zero = scale < 1e-12
        if zero.any():
            names = [n for n, z in zip(self.feature_names, zero) if z]
            logger.warning("zero-variance features %s — scale pinned to 1.0", names)
            scale[zero] = 1.0
        self.scale_ = scale
        logger.info("fitted scaler on %d rows, %d features", int(valid.sum()), x.shape[1])
        return self

    def transform(self, x: np.ndarray, feature_names: list[str] | None = None) -> np.ndarray:
        """Apply the FITTED statistics. Never recomputes anything."""
        if self.mean_ is None or self.scale_ is None:
            raise RuntimeError(
                "scaler is not fitted and must not be fitted at inference time — "
                "load the artifact written during training")
        if feature_names is not None and list(feature_names) != self.feature_names:
            raise ValueError(
                f"feature order mismatch:\n  expected {self.feature_names}\n"
                f"  received {list(feature_names)}")
        return ((x - self.mean_) / self.scale_).astype("float32")

    # ---- artifact I/O -----------------------------------------------------
    def save(self, path: str) -> str:
        blob = {"feature_names": self.feature_names,
                "mean": self.mean_.tolist(), "scale": self.scale_.tolist()}
        payload = json.dumps(blob, sort_keys=True)
        with open(path, "w") as fh:
            fh.write(payload)
        digest = hashlib.sha256(payload.encode()).hexdigest()[:16]
        logger.info("wrote scaler %s (sha256:%s)", path, digest)
        return digest

    @classmethod
    def load(cls, path: str, expect_digest: str | None = None) -> "BandScaler":
        with open(path) as fh:
            payload = fh.read()
        digest = hashlib.sha256(payload.encode()).hexdigest()[:16]
        if expect_digest and digest != expect_digest:
            raise ValueError(
                f"scaler artifact changed: expected sha256:{expect_digest}, got {digest}")
        blob = json.loads(payload)
        scaler = cls(feature_names=blob["feature_names"])
        scaler.mean_ = np.asarray(blob["mean"], dtype="float64")
        scaler.scale_ = np.asarray(blob["scale"], dtype="float64")
        return scaler


def scaling_health(scaled: np.ndarray, feature_names: list[str]) -> dict:
    """Post-scaling moments to log per inference batch."""
    finite = np.isfinite(scaled).all(axis=1)
    sub = scaled[finite]
    return {name: {"mean": round(float(sub[:, i].mean()), 4),
                   "std": round(float(sub[:, i].std()), 4)}
            for i, name in enumerate(feature_names)}

Step-by-Step Walkthrough

Step 1 — Fit on the training split, with nodata masked

FEATURES = ["blue", "green", "red", "nir", "swir1", "elevation", "slope_deg"]

x_train = train_frame[FEATURES].to_numpy("float64")
scaler = BandScaler(FEATURES).fit(x_train)
digest = scaler.save("artifacts/band_scaler.json")

Step 2 — Record the digest in the model metadata

import json

with open("artifacts/model_card.json", "w") as fh:
    json.dump({"model": "landcover_v7", "scaler": "band_scaler.json",
               "scaler_sha256": digest, "features": FEATURES}, fh)

Binding the scaler’s hash to the model is what turns “we think these go together” into a checkable fact. It is the same idea as pinning a model version once per run in orchestrating geospatial ML pipelines.

Step 3 — Load, never re-fit, at inference

card = json.load(open("artifacts/model_card.json"))
scaler = BandScaler.load("artifacts/band_scaler.json", expect_digest=card["scaler_sha256"])

x_tile = tile_frame[card["features"]].to_numpy("float64")
x_scaled = scaler.transform(x_tile, feature_names=card["features"])

If band_scaler.json is missing, this raises. That is intentional: a pipeline that quietly falls back to fitting produces plausible predictions that are wrong, which is far more expensive than a failed run.

Step 4 — Log the health check every batch

health = scaling_health(x_scaled, card["features"])
logger.info("post-scaling moments: %s", health)

for name, m in health.items():
    if abs(m["mean"]) < 1e-9 and abs(m["std"] - 1.0) < 1e-9:
        logger.error("feature %s is exactly standardised on this batch — "
                     "the scaler was almost certainly re-fitted", name)

Exact zero mean and unit variance on a production batch is not a sign of health; it is the fingerprint of a re-fit.

The scaler is an artifact with a hash, not a step to repeat The training run fits the scaler, writes its statistics to a file, and records the file's hash in the model card. The inference run loads the model card, loads the statistics, verifies the hash matches, and only then transforms the batch. If the artifact is missing the run fails rather than re-fitting. Bind the statistics to the model, then verify the binding training run fit on train split only band_scaler.json mean, scale, feature order model card records sha256 inference run tile arrives load + verify hash no fitting allowed transform then log the moments artifact missing → raise, do not fall back to fitting

Verification

import numpy as np
import pytest


def test_transform_is_batch_independent():
    """The same raw row must scale identically regardless of what it is batched with."""
    rng = np.random.default_rng(0)
    train = rng.normal(0.3, 0.1, size=(5000, 3))
    scaler = BandScaler(["a", "b", "c"]).fit(train)

    row = np.array([[0.42, 0.31, 0.27]])
    alone = scaler.transform(row)
    with_others = scaler.transform(np.vstack([row, rng.normal(0.9, 0.05, (500, 3))]))[0:1]
    assert np.allclose(alone, with_others)


def test_inference_refuses_to_fit():
    scaler = BandScaler(["a", "b"])
    with pytest.raises(RuntimeError, match="must not be fitted at inference"):
        scaler.transform(np.zeros((4, 2)))


def test_feature_order_mismatch_is_caught():
    scaler = BandScaler(["red", "nir"]).fit(np.random.rand(500, 2))
    with pytest.raises(ValueError, match="feature order mismatch"):
        scaler.transform(np.random.rand(4, 2), feature_names=["nir", "red"])


def test_artifact_hash_is_enforced(tmp_path):
    path = tmp_path / "s.json"
    scaler = BandScaler(["a"]).fit(np.random.rand(500, 1))
    digest = scaler.save(str(path))
    path.write_text(path.read_text().replace("]}", ", 0.0]}"))   # tamper
    with pytest.raises(ValueError, match="scaler artifact changed"):
        BandScaler.load(str(path), expect_digest=digest)

The first test is the one that would have caught every re-fitting incident I have seen: it states the property that matters — batch independence — rather than testing the arithmetic.

A perfectly flat post-scaling mean is a bug signature A time series of the logged post-scaling mean per inference batch. For most batches the value fluctuates gently around zero, reflecting real distributional variation between tiles. Over one stretch the value is pinned at exactly zero, which indicates the scaler was re-fitted on each of those batches rather than loaded. Logged post-scaling mean, one point per inference batch 0 exactly 0.0000 — re-fitted healthy: small drift as tiles differ healthy again after the fix Real drift wanders. A bug is suspiciously perfect.

FAQ

Why can’t I just re-fit the scaler on each inference batch?

Because the model learned a mapping from scaled values to targets, and re-fitting changes what a scaled value means. A homogeneous forest tile has its mean mapped to zero, so every pixel reports as average; the same pixel in a mixed tile reports as bright. Predictions then depend on batching, which is not a property any deployment wants.

Should the scaler live inside the model artifact or beside it?

Inside, when the runtime allows — a scikit-learn Pipeline or an ONNX graph with the scaling baked in cannot be deployed without its statistics. When the runtime needs a bare tensor model, ship a versioned sidecar and verify its hash at load, as above.

How do I detect skew that has already happened?

Log post-scaling mean and standard deviation per feature for every batch. Correct behaviour is mild wander around zero and one. A batch whose scaled mean is exactly zero to machine precision was re-fitted.


Part of: Feature Scaling for Geospatial Inputs Part of: Spatial Feature Engineering for Machine Learning