🤖 machine learning

Training-Inference Kernel Contracts: Bounding Divergence in Post-Training and Deployment

This paper proposes a "kernel contracts" framework to formally specify and bound the distributional divergence between training and inference kernels in post-training pipelines, deriving theoretical bounds on policy-gradient bias and outlining a structured deployment pipeline, while noting that it presents a conceptual framework without production-scale empirical validation.

Original authors: Bruce Changlong Xu, Lan Wu

Published 2026-06-09

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Bruce Changlong Xu, Lan Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant chef (the AI model) who has spent years learning to cook in a high-end, perfectly calibrated test kitchen. In this kitchen, they use precise digital scales, fresh ingredients, and a slow, careful cooking process to ensure every dish is perfect. This is the Training Kitchen.

Now, imagine you want to serve this chef's recipes to thousands of hungry customers at a busy food truck. To keep up with demand, you switch to a different setup: you use pre-measured spice packets, a faster (but slightly less precise) grill, and a system where orders are batched together to save time. This is the Inference Kitchen.

The problem, according to this paper, is that even though the chef is the same person using the same secret recipe (the model weights), the food coming out of the food truck isn't exactly the same as the food from the test kitchen. The difference is tiny—maybe a pinch of salt here or a slightly different sear there—but over thousands of orders, these tiny differences can add up. Sometimes, a dish that was supposed to be "spicy" comes out "mild," or a safety check that worked in the kitchen fails on the truck.

The paper calls this gap the "Training-Inference Kernel Contract." Here is a simple breakdown of their solution:

1. The Problem: "Two Different Chefs"

Currently, when we build AI, we assume the "Training Chef" and the "Inference Chef" are doing the exact same thing. But in reality, they are using different tools and methods.

Training uses high-precision math (like a digital scale).
Inference uses fast, low-precision math (like a visual estimate) to go faster and save money.

Because they use different tools, they sometimes make different decisions. In a normal restaurant, this might just mean a soup tastes slightly different. But for AI, this can mean:

The "Reward Hack": In Reinforcement Learning (where the AI learns by trial and error), the AI might think it's doing a great job because the "fast" kitchen gave it a good score, while the "precise" kitchen would have given it a bad one. It's like a student getting an A on a practice test but failing the real exam because the grading rubric changed.
The "Safety Slip": A prompt that the model refuses to answer in the test kitchen might accidentally get answered in the food truck because the fast grill changed the flavor just enough to bypass the safety filter.

2. The Solution: The "Kernel Contract"

The authors propose a new rulebook called a Kernel Contract. Think of this not as a legal document for lawyers, but as a Quality Control Checklist that travels with the AI.

This contract says: "We know the fast kitchen (Inference) won't be 100% identical to the test kitchen (Training). That's okay. But here are the specific rules we will NOT break."

The contract has four main sections:

Numerical Rules (N): "The math can't drift more than X amount." (e.g., The spice level can't change by more than 10%).
Statistical Rules (S): "The final taste must be consistent." (e.g., 99% of the time, the dish must still be recognized as 'Spicy').
Runtime Rules (R): "It still has to be fast enough." (e.g., The food truck can't slow down just because we added a safety check).
Observability Rules (O): "We must be able to taste-test any specific order later." (If a customer complains, we must be able to replay that exact order in both kitchens to see what went wrong).

3. The "Escalation Policy" (What happens if you break the rules?)

The contract isn't just a list; it has a traffic light system:

Green (L1): "Heads up." We logged a tiny difference. Keep cooking.
Yellow (L2): "Warning." The difference is getting too big. We stop sending new orders to this kitchen and route them to a backup kitchen until we fix it.
Red (L3): "Emergency." Something is critically wrong. We immediately shut down this kitchen and switch to a known-good version.

4. The "Four-Stage Promotion" (How to test before serving)

You don't just flip a switch and send the new kitchen to the public. The paper suggests a four-step safety tunnel:

Offline CI: Run the checklist on a fixed set of test orders in the lab. If it fails, don't even leave the lab.
Shadow: Let the new kitchen cook, but serve the old kitchen's food to customers. We just watch to see if the new kitchen would have made mistakes.
Canary: Let the new kitchen serve a tiny group of real customers (like 1%). If they complain, we stop immediately.
Full: If everyone is happy, we let the new kitchen serve everyone.

5. Why This Matters for "Learning" AI (RL)

The paper makes a specific point about AI that learns by itself (Reinforcement Learning).

The Issue: When the AI learns, it takes a "snapshot" of the world using the fast kitchen, but then tries to learn from the precise kitchen. It's like trying to learn to drive a car by watching a video of a race car, but then driving a different model of car. The AI gets confused and learns the wrong lessons.
The Fix: The contract forces the AI to admit, "Hey, my fast kitchen and precise kitchen are different." It adds a "correction factor" to the learning process so the AI doesn't get tricked by the speed of the food truck.

Summary

The paper argues that we need to stop pretending the "Training AI" and the "Serving AI" are the same thing. Instead, we should treat them as two different partners who have signed a Contract. This contract explicitly states how much they are allowed to disagree, what happens if they disagree too much, and how to catch those disagreements before they ruin the customer's experience.

It's about moving from "hoping everything works" to "measuring exactly where the differences are and managing them."

Technical Summary: Training-Inference Kernel Contracts

Problem Statement

Modern post-training pipelines for Large Language Models (LLMs) treat the policy symbol $\pi_\theta$ as a single mathematical object. However, in finite precision, $\pi_\theta$ is evaluated by two distinct programs: a training kernel ( $K_{train}$ ) optimized for differentiability (using high-precision accumulators, fixed shapes, and reference attention) and an inference kernel ( $K_{inf}$ ) optimized for throughput (using low-precision operators, paged key-value caches, and dynamic attention tiles).

The paper identifies a critical "specification gap": these two kernels compute measurably different functions of the same weights $\theta$ . This divergence is not uniform; it concentrates on specific slices (e.g., long-context inputs, safety-critical prompts) that aggregate benchmarks often under-represent. The conflation of $K_{train}$ and $K_{inf}$ leads to several failure modes:

Off-policy bias in RL: Reinforcement learning pipelines (e.g., PPO, GRPO) sample rollouts from $K_{inf}$ but compute gradients using $K_{train}$ . The standard estimator assumes the importance ratio $w = \pi^{K_{train}}_\theta / \pi^{K_{inf}}_\theta$ is 1, introducing a bias that manifests as KL collapse, reward hacking, or sudden performance degradation.
Deployment Drift: Models fine-tuned in BF16 may drift significantly when cast to FP8 or INT4 at inference, particularly on rare safety prompts or structured outputs.
Non-reproducibility: Heterogeneous serving fleets (different hardware/kernel builds) produce divergent outputs for the same request, making A/B tests and safety audits non-reproducible.

Currently, no machine-readable specification exists to define what level of divergence between $K_{train}$ and $K_{inf}$ is tolerable, on which data slices, and what actions should be taken when limits are exceeded.

Methodology: The Kernel Contract

The authors propose a formal specification object called the Kernel Contract, denoted as $C = (N, S, R, \Pi)$ . This contract travels with the model checkpoint and kernel build through the CI/CD pipeline.

1. The Contract Structure

The contract consists of four families of clauses and an escalation policy:

Numerical ( $N$ ): Constraints on raw numerical disagreements, such as logit drift ( $\Delta_{logit}$ ), KL divergence, Total Variation (TV) distance, and the importance ratio $w$ for RL.
Statistical ( $S$ ): Constraints on behavioral outcomes, including top- $k$ token agreement, calibration drift ( $\Delta_{ECE}$ ), reward deltas, and structured output parseability.
Runtime ( $R$ ): Operational constraints on latency, throughput, memory usage (including KV cache), and failure rates.
Observability ( $O$ ): Requirements for trace fields, kernel fingerprints, and replay signatures to ensure violations are attributable.
Escalation Policy ( $\Pi$ ): A mapping from violation severity levels ( $L1$ : informational, $L2$ : quality-impacting, $L3$ : safety-critical) to routing actions (log, guard, fallback/rollback).

Clauses are defined with thresholds ( $\tau$ ), allowed exceedance rates ( $\rho$ ), and specific data slices ( $\Sigma$ ) to which they apply.

2. Theoretical Bounds (The Chain)

The paper derives a chain of Lipschitz-style bounds connecting low-level numerical errors to high-level reward errors:

Logit to Distribution: A bound on logit drift ( $\|\varepsilon\|_\infty$ ) propagates to a bound on Total Variation distance ( $\Delta_{TV}$ ) and KL divergence, scaled by temperature $T$ .
Distribution to Reward: For a bounded reward function $r$ , the expected reward difference is bounded by $2R \cdot \Delta_{TV}$ .
Application to RL: Under assumptions of support overlap and bounded advantage, the paper proves that a per-token bound on the importance ratio $w$ (or equivalently, the TV distance between $K_{train}$ and $K_{inf}$ policies) bounds the bias of the standard policy-gradient estimator. This reinterprets the PPO clipping parameter $\epsilon$ not just as a trust-region for policy updates, but as a budget for the train-infer kernel pair.

3. Enforcement Pipeline

The framework proposes a four-stage promotion pipeline to enforce the contract:

Offline CI: Evaluates hard numerical clauses against fixed datasets.
Shadow: Mirrors live traffic to the candidate kernel to detect workload-dependent regressions.
Canary: Routes a small fraction of live traffic to the candidate with the escalation policy active.
Full Promotion: Occurs only after sustained compliance over a stability window.
An online routing loop continuously monitors contract health ( $h_t$ ) and switches traffic to guarded or safe kernels based on real-time violation levels.

Key Contributions

The paper makes three primary contributions:

Conceptual: It identifies the correct unit of specification in modern LLM pipelines not as the checkpoint or the engine alone, but as the pair $(K_{train}, K_{inf})$ along with the specific data slices where they must agree.
Technical: It provides a chain of bounds allowing engineers to specify constraints at the measurable numerical level (logits) and inherit guarantees at the behavioral level (rewards).
Applied: It demonstrates that standard RL stability heuristics (like PPO clipping) can be formally interpreted as contract clauses bounding kernel divergence, and it outlines a minimal YAML DSL for implementing these artifacts.

Results and Validation

The paper explicitly states that it does not report empirical validation at production scale and does not claim kernel divergence is the sole cause of the failure modes discussed. Instead, it presents a framework and vocabulary paper.

The "results" are theoretical derivations and a proposed experimental protocol (Section 8) for future validation:

Theoretical: Proofs showing that logit drift bounds imply reward drift bounds and that kernel mismatch introduces a calculable bias in policy gradients.
Proposed Experiments: The authors outline five falsifiable experiments (E1–E5) to test the framework, including mapping numerical drift baselines, testing hardware sensitivity, validating the RL bias bound, exploring contract-aware distillation, and simulating self-healing routing under fault injection.

Significance and Scope

The paper positions itself as a foundational specification layer that bridges several parallel literatures: RL post-training, high-throughput serving, formal verification, and software contracts.

Modesty: The authors are explicit that this is a "framework and vocabulary" paper. They do not claim to have solved the problem of kernel divergence empirically but rather to have provided the "missing index" on $\pi_\theta$ and the vocabulary to discuss it.
Future Directions: The paper outlines four research directions: dynamic contracts (context-aware thresholds), contract-aware distillation (training models to satisfy contracts), formal verification of clause subsets, and self-healing routing.
Limitations: The framework adds engineering overhead (computation, storage, observability) and carries risks of Goodhart's law (optimizing proxy clauses). It acknowledges that current formal verification tools cannot yet scale to full production transformers, relying instead on loose bounds and empirical testing.

In summary, the paper argues that treating the training and inference kernels as distinct, contract-bound entities is necessary to move from "folk knowledge" about implementation fragility to auditable, reproducible engineering practices in LLM deployment.