ML Engineering in the LLM Era

Assumes you know from ML-101

This lesson builds on ML-101 Lesson 16 (MLOps & Deployment) for the basic vocabulary of training-serving pipelines, on ML-101 Lesson 9 (Model Evaluation) for the precision/recall trade-off that the retrieval-augmented generation objective reduces to, and on ML-101 Lesson 11 (Gradient Descent & Optimization) for the optimisation machinery behind fine-tuning. We will not re-derive gradient descent or define what a CI/CD pipeline is.

The reader is also expected to know ML-101 Lesson 10 (Overfitting, Bias & Variance) because every reliability story in MLOps is ultimately a story about generalisation: training metrics that fail to predict held-out metrics are a bias problem, and held-out metrics that fail to predict deployed metrics are a variance problem. The LLM era adds a third axis — drift — and the diagnostic machinery that works for the first two axes must be extended, not replaced, to handle drift.

Learning Objectives

1. Derive the retrieval-augmented generation (RAG) objective from the chain rule of conditional probability, and explain why the retrieval step is necessary rather than convenient, identifying the assumption whose violation makes the LLM's parametric knowledge trustworthy enough to skip retrieval.
2. Compute the cost of a RAG pipeline — vector-store indexing, retrieval latency, context-window consumption, and per-token generation cost — and predict which cost dominates at a given scale from first principles rather than from benchmarks.
3. State the population-stability index and the Wasserstein-based drift detector, derive the threshold that makes the detector's false-positive rate acceptable, and explain why neither detector can identify drift without a reference window.
4. Design a continuous-training pipeline that handles catastrophic forgetting via experience replay and Elastic Weight Consolidation (EWC), and justify the EWC penalty from the Laplace approximation of the Bayesian posterior over the previous task.
5. Justify the LLM-specific deployment pattern of shadow deployment for prompt changes, and explain why A/B testing on the model output alone is insufficient when the loss surface has been replaced by a human-rater surface.

RAG: the retrieval objective

Retrieval-augmented generation (RAG) augments a language model's generation step with a retrieval step that fetches documents from an external corpus at inference time. The retrieval step is what lets the model answer questions about events that occurred after its training cut-off, cite its sources, and avoid hallucinating on facts that are in the corpus but not in its weights. The objective that RAG optimises is implicit, and deriving it shows why the retrieval step is necessary rather than convenient.

The probability of an answer $y$ given a query $x$ decomposes by the chain rule of conditional probability:

The outer sum is over a (typically large) corpus $\mathcal{D}$, and $p(d \mid x)$ is the retriever's score for document $d$ given query $x$. The vanilla language model without retrieval corresponds to marginalising over a corpus that lives inside the model's weights — the implicit "parametric memory" — and the marginalisation is performed by the attention mechanism during the forward computation. The two approaches are mathematically equivalent in form; they differ in where the corpus lives and in what cost function governs the choice of marginalisation.

The retrieval step is necessary rather than convenient when the parametric memory is incomplete. The implicit assumption is that the parametric memory is not a faithful approximation of the marginal $\sum_d p(y \mid x, d)\,p(d \mid x)$ over the documents we actually care about. Two failure modes violate this assumption in practice:

1. Recency: the parametric memory has a training cut-off; documents written after the cut-off are not in the weights and the model's output cannot reflect them. Retrieval closes the cut-off.
2. Specificity: the parametric memory has finite capacity; documents that are rare in the training distribution are not well-modelled by the weights. Retrieval closes the long-tail.

A 101-level misconception is that RAG "makes the model smarter". RAG does not change the language model at all — it provides the model with context that was not in its weights, and the model uses its existing mechanism to condition on that context. A strong retrieval pipeline combined with a weak language model outperforms a weak retrieval pipeline combined with a strong language model on most knowledge-intensive benchmarks, because the bottleneck is rarely the language model — it is the relevance of the retrieved context.

import numpy as np

def rag_answer(query: str, retriever, generator, k: int = 5) -> str:
    """
    k is the number of documents in context. Doubling k roughly doubles
    the prompt-token cost and may push the answer out of the context
    window entirely if the documents are long.
    """
    # Score every document in the corpus under the query.
    doc_scores = retriever.score(query)            # shape (|D|,)
    top_k = np.argsort(-doc_scores)[:k]             # indices of top k
    docs = [retriever.corpus[i] for i in top_k]
    # Concatenate and generate. The generator's loss is unchanged;
    # the conditional p(y | x, d) has the same form as before, just
    # with a different x.
    prompt = format_prompt(query, docs)
    return generator.generate(prompt)

The retrieval objective also has a subtle evaluation consequence: the right metric is not BLEU or ROUGE against a reference answer, but answer faithfulness against the retrieved context. A model can produce a fluent answer that contradicts the retrieved documents, and that answer is a RAG failure even if it scores high on a language-quality metric. The standard practice is to use an LLM-as-judge to evaluate faithfulness: feed the judge the retrieved documents and the generated answer, and ask whether the answer is supported by the documents. This is a noisier signal than a learned classifier would be, but the labelled data is essentially free.

Cost accounting for an LLM serving pipeline

The total cost of serving an LLM with RAG has four components, and the relative weight of each changes with deployment scale. Let $Q$ be the queries-per-second, $D$ the corpus size in tokens, $k$ the documents-per-query, $L$ the average output length in tokens, and $C$ the per-token generation cost. The four components are:

The table is not folklore; it follows from how each component scales with the deployment parameters. The vector index uses approximate nearest-neighbour (ANN) search with sub-linear query complexity; the context-window cost is per-second tokens consumed across all in-flight requests; the generation cost is the only component that scales with the LLM's parameter count, because every token requires a forward pass through the full network.

The dominant cost at production scale is almost always generation. A 7B-parameter model serving 100 QPS with an average output of 200 tokens per query requires $100 \times 200 = 20{,}000$ generated tokens per second, which on an A100 GPU at $\sim 30$ tokens per millisecond is roughly 30% of one GPU's throughput. The retrieval step on the same hardware adds less than 5 ms per query. The lesson is that the choice of model is the dominant cost decision, not the choice of retriever, and the most powerful cost lever is reducing $L$ (output length) — for example, by switching to a smaller model or by enforcing structured outputs that are shorter than free-form prose.

A 101-level misconception is that RAG "costs almost nothing" because the documents are pre-indexed. Retrieval latency is small but not zero, and context-window consumption has a hard ceiling that RAG cannot exceed. A pipeline that retrieves 50 long documents per query will saturate the context window long before the retrieval latency becomes the bottleneck.

Drift detection: when production ≠ training

Drift is the phenomenon that the production input distribution $P_T$ diverges from the training input distribution $P_S$. It is distinct from concept drift (the labelling function changes) and from label drift (the marginal $p(y)$ changes), though all three are commonly conflated. The diagnostic machinery for input drift is a hypothesis test against a reference window of training-time inputs.

The simplest detector is the population-stability index (PSI), a discrepancy statistic on binned univariate marginals. Bin both $P_S$ and $P_T$ into $K$ bins (typically $K = 10$ or $20$), let $p_k$ and $q_k$ be the bin probabilities, and compute

This is a discretised KL divergence on the histogram. The empirical rule of thumb in production MLOps is that PSI $< 0.1$ indicates no significant drift, $0.1 \le \text{PSI} < 0.25$ indicates moderate drift, and PSI $\ge 0.25$ indicates significant drift. The thresholds are PSI-tuned to a target false-positive rate of approximately 5% on production data; they are not derivable from first principles and should be re-calibrated for each application.

The PSI is a univariate detector: it catches drift in the marginal distributions of individual features but is blind to covariate drift — a rotation of the feature space that preserves marginals but changes the joint. A richer detector uses the Wasserstein distance between the empirical distributions of $P_S$ and $P_T$:

where $\Gamma$ is the set of couplings. The 1-Wasserstein has a closed-form expression in 1D as the integral of the absolute difference of CDFs; in higher dimensions it is computed via the Sinkhorn algorithm with an entropic regulariser. The threshold for the 1-Wasserstein detector is set by bootstrap resampling of $P_S$: the 95th-percentile of the bootstrap distribution is the threshold, giving a controlled 5% false-positive rate.

The deeper problem is that neither detector can identify drift without a reference window. The training distribution is the default reference, but if the deployment has been running for months without re-training, the production distribution may have drifted from the training distribution in two ways — the change from training to current production, and the change from current production to itself over time. A useful diagnostic is to monitor the detector against a rolling window as well as the original training distribution, and alert only when both windows agree that drift has occurred.

import numpy as np
from scipy.stats import wasserstein_distance

def psi(p_ref: np.ndarray, p_cur: np.ndarray, eps: float = 1e-6) -> float:
    """Population-stability index on two equal-length histograms."""
    p = np.clip(p_ref, eps, None); q = np.clip(p_cur, eps, None)
    return float(np.sum((q - p) * np.log(q / p)))

def drift_alert(scores_ref: np.ndarray, scores_cur: np.ndarray,
                bins: int = 20, psi_threshold: float = 0.2) -> dict:
    """Compare the empirical distribution of model scores (or any 1D signal)."""
    hist_ref, edges = np.histogram(scores_ref, bins=bins, density=True)
    hist_cur, _     = np.histogram(scores_cur, bins=edges, density=True)
    return {
        "psi":           psi(hist_ref, hist_cur),
        "wasserstein_1": wasserstein_distance(scores_ref, scores_cur),
        "psi_alert":     psi(hist_ref, hist_cur) >= psi_threshold,
    }

A 101-level misconception is that drift detection is solved by "comparing the means". Comparing means catches only the first moment of a distribution; the production input distribution can shift its variance, skew, multimodality, or any higher moment without the means moving. A robust detector uses at least one divergence that is sensitive to the full distribution, not just the mean.

Continuous training and forgetting

A deployed model needs to be re-trained when drift crosses a threshold. The naive approach — fine-tune on the new data — suffers from catastrophic forgetting: the model loses competence on the old data while gaining it on the new. The LLM era adds two standard mitigations, both of which derive from a Bayesian framing of the continual-learning problem.

Experience replay stores a small buffer of representative samples from past tasks (or generates them from the model's own parametric memory) and mixes them into the new-task training set. The mixed loss is

where $\mathcal{B}_{\text{old}}$ is the replay buffer. The replay loss is the simplest possible regulariser against forgetting, and it is effective when the buffer is representative of the old distribution.

Elastic Weight Consolidation (EWC) uses a quadratic penalty on the parameter movement, weighted by the Fisher information of the old task at the previous optimum $\theta_S^{\star}$:

The Fisher diagonal $F_{ii}$ measures how sensitive the old-task loss is to $\theta_i$, and the penalty is larger for parameters that mattered more for the old task. The derivation is the Laplace approximation of the Bayesian posterior: the posterior over $\theta$ given the old task is approximately Gaussian around $\theta_S^{\star}$ with covariance equal to the inverse Fisher, and the new-task posterior is approximately Gaussian around its own optimum; the EWC loss is the negative log of the product of these two Gaussians, evaluated at $\theta$.

def ewc_penalty(model, fisher: dict, theta_star: dict, lam: float = 1e4):
    """Sum_i (lam/2) F_ii (theta_i - theta_S_i)^2 — second term of EWC loss."""
    total = 0.0
    for name, p in model.named_parameters():
        if name in fisher:
            total += (fisher[name] * (p - theta_star[name]).pow(2)).sum()
    return 0.5 * lam * total

The Fisher information requires a pass over the old-task data at $\theta_S^{\star}$, which means it must be computed before re-training begins and stored. It is a per-parameter, per-task quantity; in a multi-task continual learning setting the Fisher accumulates across tasks. The Laplace-approximation assumption is that the loss surface near $\theta_S^{\star}$ is approximately quadratic — true for well-trained models far from saddle points but violated for models that are still in the early-training regime.

Deployment patterns for prompts and models

The deployment pattern that distinguishes the LLM era from classical ML is that prompts are code: they are versioned, reviewed, tested, and rolled back, and a prompt change has the same production-impact magnitude as a model change. The standard deployment pattern for a prompt change is shadow deployment: the new prompt is run alongside the current one for a fraction of traffic, and its outputs are logged for offline evaluation before the prompt is promoted.

Shadow deployment differs from A/B testing in one crucial way: the new prompt's outputs are not shown to users. They are evaluated against an automated metric (LLM-as-judge, a regression test suite, or a domain-specific evaluator) and compared to the old prompt's outputs on the same inputs. A/B testing is required once the new prompt is promoted, because the end-user response to a changed output is a behavioural signal (click-through, conversion, retention) that no offline metric captures.

The deeper reason is that the loss surface has been replaced by a human-rater surface. Classical A/B testing relies on the expectation that small input changes produce small output changes — the smoothness assumption that makes gradient-based optimisation work in the first place. A prompt change can flip a model's output from "yes, I will help with that" to "I cannot help with that" with no intermediate state, which violates smoothness and makes the A/B test's variance dominate its mean. Shadow deployment replaces the human-rater signal with an automated one before the A/B test, so the A/B test is comparing two outputs that are both already known to be acceptable.

A useful deployment checklist for prompt changes:

1. Regression suite: a fixed set of inputs with expected properties in the output. Run the new prompt through the suite and check for property violations.
2. Shadow deployment: 1–5% of production traffic, parallel to the current prompt, for at least 24 hours.
3. Canary: 1% of production traffic on the new prompt with the user seeing the new prompt's output, with the ability to roll back in under five minutes.
4. A/B test: 50/50 traffic split, with a primary metric (e.g. task completion rate) and a guardrail metric (e.g. safety-violation rate).
5. Full rollout: 100% traffic on the new prompt.

The five-step sequence is the LLM-era analogue of the ML-101 deployment checklist, and skipping any step (especially shadow deployment) is the most common cause of prompt-rollback incidents.

Key Takeaways

• RAG is a Bayesian marginalisation over an external corpus: the retrieval step is necessary when the parametric memory is incomplete or out of date, and the standard failure modes are recency and specificity. The bottleneck on knowledge-intensive tasks is rarely the language model — it is the relevance of the retrieved context.
• Cost accounting for an LLM pipeline has four components, and the dominant cost at production scale is almost always generation, not retrieval. Reducing output length is the single most powerful cost lever.
• Drift detection requires a reference window and a divergence statistic; the PSI is a discretised KL on binned histograms and is appropriate for univariate marginals, while the 1-Wasserstein is appropriate for the full empirical distribution. Neither can identify drift without a reference, and the rolling-window diagnostic catches the second-order drift that single-reference monitoring misses.
• Continuous training without forgetting uses experience replay or EWC; the EWC penalty is the Laplace approximation of the Bayesian posterior over the previous task, weighted by the Fisher information at the previous optimum.
• Prompt changes are code changes and require a five-step deployment pattern (regression suite, shadow, canary, A/B, full rollout). Shadow deployment is the LLM-era analogue of the unit-test stage and is non-negotiable for any prompt that ships to production.