How We Fixed Training-Serving Skew in Our ML Models

Ask a data science team how long they take to go from “we have a model idea” to “the model scores live traffic.” The honest answer rarely involves the model. You fit a gradient-boosted tree or a GLM in an afternoon. A few lines of a well-documented library, a grid search, done.

The weeks go somewhere else. They go into the work around one call to .fit(): wiring up data, deriving the same features twice, writing a second “production” copy of the preprocessing, proving to a reviewer the copy matches the original, and then hoping nothing drifts after launch.

Most applied ML engineering effort lands in this gap, between a model which works in a notebook and a model safe to run unattended on live data. You don’t have to live with the gap. This post covers why the gap exists, why teams get the gap wrong so often, and what we learned building AXIOM, PhonePe’s internal framework for this middle layer.

The lessons apply whether you use Spark, pandas, or something else. Every claim below points to something we built.

Where Data Scientists Actually Add Value

A model is a function. Features go in, a prediction comes out. The interesting work for a data scientist sits in three decisions: 

1. Choosing the features that actually matter.
2. Picking the right algorithm for the job (classification, regression, or ranking).
3. Making sure the function holds up on new data.

These decisions need human judgment. No framework should try to automate them.

Everything else repeats. Reading data, joining tables correctly, handling nulls, encoding categories, splitting the same way each time, tracking what ran and why. Each step looks easy because each step is easy. The trouble starts when you have to run every step twice, identically: once in training, once when scoring new data. And at a scale where nobody checks each batch by hand.

What Goes Wrong When You Rewrite Preprocessing Code?

We kept seeing one failure pattern. Other companies see the same one. A data scientist builds a notebook, fits a model, gets a good AUC, and someone asks them to “productionize” the work. So a second engineer, or the same person under deadline, reads the notebook and rewrites the preprocessing as a “clean” inference script. Null handling, categorical encoding, feature derivations, all rewritten by hand.

Most production ML incidents start in this rewrite. The reason is simple. Two separate implementations of the same logic drift apart, even when one person writes both in the same week. Some examples:

• The training notebook computes an imputation median on the training data. The inference script recomputes the median on whatever small batch shows up. Same code, different number.
• A categorical encoder trained on known labels meets an unseen category. The rewritten version drops the row, buckets the value as “unknown,” or crashes, and nobody chose which on purpose.
• A date feature anchors on “today.” In training, “today” stays fixed at train time. At inference, “today” moves, and every recency feature shifts with no warning.
• The simplest one. Someone fixes a bug in the training script after launch and forgets the parallel inference script exists.

None of these crash. They show up weeks later as a slow accuracy drop nobody explains, because the model now scores inputs different from the ones in training. People call this training-serving skew. In almost every case, the root cause is duplicated code.

Training and serving run as separate jobs, and both use one artifact: the MOJO. Serving never rewrites preprocessing. Serving loads the exact binary from training. DriftGuard compares the distribution the artifact trained on with the live data the artifact scores. When the PSI gap grows large enough to matter, DriftGuard alerts the model owner, who reviews the drift and decides whether to retrain.

How AXIOM Keeps Training and Serving in Sync

A checklist, a wiki page, or a stricter PR review won’t solve training-serving skew. All three depend on people remembering a manual check every time, forever. People forget. The only lasting fix removes the option of running two implementations. Train and serve must run the exact same artifact.

AXIOM follows this rule. Three design choices enforce the rule:

1. Preprocessing is a typed plugin, not a script. Every transformation (null handling, encoding, imputation) implements one abstract interface with fit and transform. A declarative list in a config file calls each transformer. Nobody hand-writes transformers per use case. No “training version” or “inference version” exists. You get one class, instantiated once.
2. Training produces one self-contained, serializable artifact. AXIOM’s model layer exports an H2O MOJO (Model Object, Optimized). This binary bundles the fitted model with the transformations fitted alongside the model. Scoring runs no Python preprocessing. Scoring loads the binary from training and reuses the binary as-is. Scoring refuses any artifact training didn’t produce. No code path exists for “re-implement this part.”
3. Config is the only user input. One YAML file declares the whole pipeline: data source, feature joins, preprocessing order, hyperparameters. A typed schema checks the YAML before any Spark job starts. A bad config gives you a readable error in seconds, instead of the silent skew you find in week six.

The result: nobody needs to check by hand whether the inference code matches training code. You have one code path to review.

Why Model Reviews Got Faster

We didn’t aim for faster reviews at first, but we noticed the change quickly. When every model on a team runs through one transformer interface and one declarative schema, reviewers stop reading custom infrastructure code. No more Hive reads, null-handling branches, or encoder logic unique to each model. Reviewers read a YAML diff instead: which features changed, which hyperparameter moved, which transformer got added.

The review surface shrinks and gets easier to read. A reviewer who knows the framework evaluates a new model’s config in minutes, instead of rebuilding someone’s custom pipeline from scratch. The framework already answered the infrastructure questions once. Does the pipeline read data correctly, join features, encode categories, and avoid skew? Those answers hold for every model. What’s left for review is the modeling decision, the part where human judgment counts.

We won’t put a percentage on the time saved. We haven’t run a controlled measurement. The direction holds, though. Standardize the 90% of a pipeline which is infrastructure, and every later review covers only the 10% which is new.

Why Models Reach Production Sooner

The same choices which fix training-serving skew also shorten deployment, for a mechanical reason. Most of the old per-model “productionization” work now ships once, as a framework feature. You no longer write an inference script, wire the script to the feature store, or hook up monitoring for each model.

In our setup:

• Feature joins stay point-in-time correct by design. Training and scoring use the same connector, so nobody rewrites a join per model.
• Drift monitoring comes attached from day one. AXIOM snapshots each model’s training feature distribution and compares the snapshot with live scoring data using Population Stability Index, the standard metric in credit risk modeling. When the input distribution shifts enough, AXIOM alerts the model owner to review the drift and retrain. The owner writes no monitoring code and keeps the final call on retraining.
• AXIOM tracks every run by default. One versioned tracking run holds the config, label logic, model artifact, and metrics. You never have to reconstruct which pipeline produced a model.

None of these ideas are new on their own. Point-in-time joins, PSI drift checks, and experiment tracking all have long track records. We didn’t invent them. We made them the default path, with zero extra effort, instead of a checklist item a team skips under deadline.

Why Do Teams Keep Repeating This Mistake?

If none of this is new, why do teams keep rebuilding the same pieces badly? A few honest reasons, true for any org scaling applied ML:

• Early models are simple. One dataset, one owner, manual oversight. Doing everything twice, carefully, works at a small scale. The problem shows up once you have too many models to audit by hand, and by then a dozen pipelines already follow the ad-hoc pattern.
• Skew stays quiet. A duplicated preprocessing bug doesn’t fail loudly. A business metric slides slowly, so teams patch the symptom and leave the structural cause alone.
• Standardization looks like a slowdown, since teams trade some flexibility for a shared interface. We saw the opposite. Once the shared interface existed, new models shipped faster, because nobody had to redo the infrastructure work. You pay to build the framework once. Every model after the first gets the payoff.

Key Takeaways

Fitting the curve was never the hard part of production ML. The hard part is guaranteeing your model sees in production exactly what the model saw in training, every time, with no human re-checking by hand. Fix the structure with one artifact, one config, and one code path. Faster reviews and shorter deployments follow directly, without a separate initiative to fund.

Typed configs, single-artifact serving, and standard drift metrics have each existed for decades. The gain comes from refusing to let any model skip them.

This post describes engineering patterns from an internal ML framework at PhonePe. The core lesson applies to any team with more than a handful of models: enforce one code path from training to serving.