Systems Specialization · CPU ReferenceCPU reference; target qualification separateLesson 1 of 8 in Phase 15

Phase 15 / Deployment, interoperability & edge AI

Adapt a pretrained model and choose a post-training objective

Revision v0.7.0 · 2026-10-20

Lesson · about 130 minutes · CPU practice

Pretrain and reload an artifact with recorded data and checkpoint hashes Compare supervised adaptation and teacher distillation from the same starting weights

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Data & checkpoint recovery; Performance diagnosis
Hardware: CPU

Download notebook

Other downloads

Python worked examples Lesson text

Setup help: Python, notebook, or Cloud TPU VM

Use the tested CPU packages (the recorded environment uses Python 3.14.3). Other Python versions and accelerator environments require separate checks. The optional Cloud TPU specialization explains the target-specific execution boundaries. Logical CPU devices practice partitioning; they do not simulate TPU performance.

Download the repository ZIP or use your existing checkout. Extract it and open a terminal in the folder containing requirements-cpu.txt. The beginner ZIP is separate; these lessons use the full repository folder.

macOS / Linux terminal · repository folder

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-cpu.txt

PowerShell · repository folder

py -3 -m venv .venv
.\.venv\Scripts\python.exe -m pip install -r requirements-cpu.txt

After activation, run the complete worked example from that same folder:

macOS / Linux terminal · activated environment

python phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/code/main.py

For an incremental build, create main.py in your own practice folder, append each build block in order, and run python main.py with this environment active. Keep that file separate from the downloaded reference solution.

Download this lesson’s notebook. Open it with a Python kernel using the course environment. Run cells from the top after restarting the kernel; experiments depend on earlier definitions.

For local Jupyter, first install the course requirements using the Local Python steps, then use your environment’s Python to start Jupyter:

macOS / Linux terminal · activated course environment

python -m pip install jupyterlab
python -m jupyter lab

PowerShell command · repository folder

.\.venv\Scripts\python.exe -m pip install jupyterlab
.\.venv\Scripts\python.exe -m jupyter lab

Select the downloaded deployment-02.ipynb file in JupyterLab. The notebook contains reference solutions; make your prediction and attempt each practice task before opening them. Restart the kernel before changing logical device setup.

Check: Both adapted branches reduce the common held-out target cross-entropy relative to the reloaded source model. The printed table also reports probability error and source-task retention; inspect all three quantities.

Revisit the guided CPU setup · Open 19: TPU Setup. Package instructions are CPU instructions; Windows and accelerator environments need their own validation.

Phase 15.01 Stage 1 of 5 · Understand

1 Understand 2 Build & Run 3 Experiment 4 Exercise & Check 5 Evidence

What you will be able to do

Pretrain and reload an artifact with recorded data and checkpoint hashes
Compare supervised adaptation and teacher distillation from the same starting weights
Verify both objective gradients independently
Report target-task gains and source-task retention without leaking held-out data

The problem

A model already learned one task. You now have a small adaptation set and a teacher that supplies probabilities. Should you optimize the observed labels or imitate the teacher? We will actually pretrain, save and reload a tiny model, adapt it both ways, and compare the resulting predictions on an untouched evaluation split. The model is deliberately small so every objective and derivative can be inspected.

The idea

Adaptation starts with the behavior you want to change and the evidence you have. Supervised labels and teacher distributions supply different targets. Their objectives can use the same model while asking it to learn different things.

Choose the target before choosing the adaptation method

A hard target identifies one class. A teacher distribution can assign mass to several alternatives, preserving information about relative preferences. Matching that distribution does not guarantee the teacher is correct for your deployment population.

Keep the evaluation target fixed when comparing adaptation methods. A lower training loss under one objective cannot be compared directly with a differently defined loss under another objective.

The held-out cross-entropy bars compare the lesson's two procedures under a common scoring rule. Explain that rule and inspect changed examples before choosing a method. This experiment is a bounded adaptation comparison; reward modeling and policy optimization require their own distinct data and update contracts.

Define the tasks before choosing an objective

The input has two measured features plus a constant intercept feature. The source task uses one synthetic probability rule; the target task uses another. Source training sees $256$ rows. Adaptation sees only $48$ different rows. The final target report uses $512$ separately generated rows. These are disjoint generated splits with separate keys, not a real benchmark. Binary adaptation labels are sampled from the target probabilities; the teacher supplies the probabilities themselves. That privileged teacher is declared explicitly rather than presented as a generally superior model.

Understand hard labels and soft targets

For logit $z=x^\top w$ , cross-entropy is $\log(1+e^z)-tz$ , where $t$ is either a binary observed label or a teacher probability between zero and one. Differentiating gives $(\operatorname{sigmoid}(z)-t)x$ . Averaging over rows gives the gradient in our host reference. A soft target contains uncertainty that a single sampled binary label does not. If the teacher is wrong, faithfully matching that uncertainty can also preserve its mistakes.

L(w)=\frac1N\sum_i[\operatorname{softplus}(x_i^\top w)-t_i x_i^\top w],\qquad \nabla L=\frac1N X^\top(\operatorname{sigmoid}(Xw)-t)

Make pretrained mean previously trained

We start from zero, optimize the source objective for $300$ updates and check that the source loss falls. We save those resulting weights, hash the checkpoint bytes and record the source-array hash, objective, dtype, version and training recipe. Adaptation uses only the reloaded artifact. This establishes a small reproducibly pretrained fixture; it is not a foundation model, a downloaded checkpoint or a random frozen feature masquerading as pretraining. The artifact omits optimizer state, so no exact-resume claim is made.

Choose a fair comparison

Both adaptation branches receive identical initial weights, adaptation inputs, update count and rate. Their training losses are not directly a competition because their targets differ. We evaluate both on one common target cross-entropy and probability squared error, using held-out generating probabilities unavailable to training. We also evaluate source loss to expose forgetting. Do not select a learning rate using this final report; add a separate validation split for model selection. One seeded synthetic result is separate from which objective will win on another task.

Locate preference and reinforcement objectives

A pairwise preference objective asks which of two responses is preferred; its supervision consists of prompt, chosen response and rejected response, with a defined sampling and reference policy. A teacher cross-entropy uses distributions over labels or tokens and does not require preference pairs. An RL objective uses a reward and a policy whose generated actions affect the training samples. This lab executes supervised adaptation and distillation only. Renaming teacher probabilities “rewards” would not create an RL experiment.

Carry the contract to a real pretrained model

Choose an actual model checkpoint under its license, record repository and exact revision, tokenizer or feature preprocessing, architecture configuration, file hashes and base precision. First verify a golden inference example before changing weights. Build disjoint training, validation and final evaluation sets, with no overlap from teacher-generation or calibration data. Start with a small controlled adapter or declared trainable parameter subset; record what remains frozen and save adapter/base identities together. Define the task and source-retention criteria before optimization. Recheck export, precision and target-runtime parity after adaptation. This extension is an execution plan, not a claim that a large external model ran in this CPU lesson.

Turn task tradeoffs into a release decision

A lower target loss can be useful while still causing an unacceptable regression on an older task. Decide how much source-task degradation is allowed before inspecting the candidate results. Let $L_s$ denote source evaluation loss and let $\Delta_s=L_s(\mathrm{candidate})-L_s(\mathrm{base})$ . Positive $\Delta_s$ means worse retention. A retention budget is a product decision expressed in the same loss units, not a universal constant.

In this synthetic exercise the source evaluation reuses the source inputs, so it is only a retention proxy. A real release needs a separate held-out source set as well as a held-out target set. Record every candidate, including candidates rejected by the gate. If none satisfies the predeclared budget, keep the base or change the training recipe; do not quietly relax the budget after seeing the answer.

\mathrm{accept}(m)=\left[L_t(m)<L_t(m_0)\right]\land\left[L_s(m)-L_s(m_0)\leq\delta\right]

Two adaptation objectives evaluated on the same held-out target

Predict before looking: Will improving the new task necessarily preserve performance on the source task?

Read the figure

The horizontal categories are the reloaded pretrained model, supervised adaptation and teacher distillation. The vertical axis is mean cross-entropy in nats. One series evaluates the same held-out target probabilities for every model; the other evaluates source-task probabilities. Lower is better within each series. The two adapted models improve target-task loss while their source-task loss increases. These bars are evaluations of the resulting weights, not the branches’ incomparable training objectives. In this seeded run the held-out target bars are approximately $0.978,0.555,0.524$ , while the source bars are $0.545,0.790,1.005$ , in the same model order. Teacher adaptation gives the lowest target loss here but also the largest source-retention loss.

Connect it to the computation

The teacher exactly matches the declared synthetic target generator, making this an especially favorable distillation case. The held-out comparison demonstrates an objective tradeoff on this fixture. It neither establishes that a real teacher is correct nor that one adaptation recipe is generally best. The failure experiment makes the teacher wrong and rechecks the target criterion.

Open full-size figure Download data and provenance

The notebook contains the data experiment, plotting code, and recorded figure output. Rerun it to test your changes.

Where to put & run this code: Work in phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/code/main.py (or downloaded exercises/deployment-02.py / exercises/deployment-02.ipynb) and execute from the repository root with python3 phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/code/main.py.

1 Pretrain a small model and retain its source data identity

Create main.py in the CPU environment. These are explicit synthetic classification tasks; the source model is genuinely optimized before it is saved.

main.py · append build step 1

# Step 1 — 1. Pretrain a small model and retain its source data identity: The same stable cross-entropy accepts binary labels or teacher...
import hashlib
import json
import tempfile
from pathlib import Path
import numpy as np
import jax
import jax.numpy as jnp

def design(key,n):
    # Sample deterministic random values into `features` using an explicit PRNG key.
    features=jax.random.normal(key,(n,2))
    return jnp.concatenate([features,jnp.ones((n,1))],axis=1)
# Create or split explicit PRNG key(s) (`source_X`) for reproducible randomness.
source_X=design(jax.random.key(10),256)
source_targets=jax.nn.sigmoid(source_X@jnp.array([1.2,-.8,.2]))
# Create or split explicit PRNG key(s) (`adapt_X`) for reproducible randomness.
adapt_X=design(jax.random.key(11),48)
# Create or split explicit PRNG key(s) (`held_X`) for reproducible randomness.
held_X=design(jax.random.key(12),512)
teacher=jnp.array([.6,1.4,-.3])
soft_targets=jax.nn.sigmoid(adapt_X@teacher)
# Create or split explicit PRNG key(s) (`hard_targets`) for reproducible randomness.
hard_targets=jax.random.bernoulli(jax.random.key(13),soft_targets).astype(jnp.float32)
held_prob=jax.nn.sigmoid(held_X@teacher)
def objective(w,X,targets):
    logits=X@w
    return jnp.mean(jnp.logaddexp(0.,logits)-targets*logits)
def train(initial,X,targets,steps=300,rate=.15):
    def step(w,_):
        # Evaluate both scalar loss and parameter gradients in one pass (`(loss, grad)`).
        loss,grad=jax.value_and_grad(objective)(w,X,targets)
        return w-rate*grad,loss
    return jax.lax.scan(step,initial,None,length=steps)
# Allocate initialized array `(base, pretrain_history)` with the specified shape and dtype.
base,pretrain_history=train(jnp.zeros(3),source_X,source_targets)
assert pretrain_history[-1]<pretrain_history[0]-.1
# Convert `source_hash` to a host NumPy array for inspection or verification.
source_hash=hashlib.sha256(np.asarray(source_X).tobytes()+np.asarray(source_targets).tobytes()).hexdigest()

Code building blocks in this step

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
jnp.zeros / jnp.ones(shape, dtype=...) — Allocates a tensor of the given shape initialized with constants.
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jax.value_and_grad(loss_fn)(params, ...) — Evaluates both the scalar loss and its gradient PyTree (loss_val, grads) in a single forward+backward pass.

The same stable cross-entropy accepts binary labels or teacher probabilities. Those targets express different supervision, even though the algebra is shared.

2 Save and reload before adapting

Append this artifact boundary. The manifest names the model, synthetic source split, objective and training recipe; adaptation starts from reloaded weights.

main.py · append build step 2

# Step 2 — 2. Save and reload before adapting: Both branches begin with identical pretrained parameters and a...
# Create an isolated temporary directory to run and inspect artifacts safely:
with tempfile.TemporaryDirectory() as directory:
    checkpoint=Path(directory)/'pretrained.npz'
    # Convert `` to a host NumPy array for inspection or verification.
    np.savez(checkpoint,weights=np.asarray(base))
    # Compute deterministic cryptographic digest `digest` for provenance verification.
    digest=hashlib.sha256(checkpoint.read_bytes()).hexdigest()
    # Compute deterministic cryptographic digest `manifest` for provenance verification.
    manifest={'model':'three-parameter logistic synthetic fixture','source_data_sha256':source_hash,
              'weights_sha256':digest,'source_seed':10,'pretrain_steps':300,'learning_rate':.15,
              'objective':'binary cross entropy with synthetic source probabilities',
              'jax':jax.__version__,'dtype':'float32','optimizer_state_included':False}
    manifest_path=Path(directory)/'manifest.json'
    # Write the serialized artifact payload to disk.
    manifest_path.write_text(json.dumps(manifest))
    recorded=json.loads(manifest_path.read_text())
    assert hashlib.sha256(checkpoint.read_bytes()).hexdigest()==recorded['weights_sha256']
    with np.load(checkpoint,allow_pickle=False) as saved:
        # Create device-backed JAX array `restored`.
        restored=jnp.asarray(saved['weights'])
# Execute `np.testing.assert_array_equal(restored,base)`
np.testing.assert_array_equal(restored,base)
frozen=np.array(restored,copy=True)
sft,sft_history=train(restored,adapt_X,hard_targets,steps=200)
distilled,distill_history=train(restored,adapt_X,soft_targets,steps=200)
# Execute `np.testing.assert_array_equal(restored,frozen)`
np.testing.assert_array_equal(restored,frozen)

Code building blocks in this step

tempfile.TemporaryDirectory(...) — Call tempfile.TemporaryDirectory with your updated parameters or inputs from this lesson's workspace.
Path(...) — Call Path with your updated parameters or inputs from this lesson's workspace.
assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

Both branches begin with identical pretrained parameters and a fresh stateless SGD update rule. This is adaptation, not exact optimizer-state resumption.

3 Compare on a held-out criterion shared by both objectives

Append the independent derivative and final evaluation. Neither optimization branch receives held_X or held_prob.

main.py · append build step 3

# Step 3 — 3. Compare on a held-out criterion shared by both objectives: The teacher is the known generating rule in this fixture, so...
w_check=jnp.array([.2,-.3,.1])
# Convert `host_X` to a host NumPy array for inspection or verification.
host_X=np.asarray(adapt_X,dtype=np.float64)
# Convert `host_w` to a host NumPy array for inspection or verification.
host_w=np.asarray(w_check,dtype=np.float64)
host_p=1/(1+np.exp(-(host_X@host_w)))
for targets in (hard_targets,soft_targets):
    # Convert `expected` to a host NumPy array for inspection or verification.
    expected=host_X.T@(host_p-np.asarray(targets))/len(host_X)
    # Differentiate the objective to obtain gradients ``.
    np.testing.assert_allclose(jax.grad(objective)(w_check,adapt_X,targets),expected,rtol=1e-5,atol=1e-6)
metrics={}
for name,weights in [('pretrained',restored),('supervised',sft),('teacher',distilled)]:
    metrics[name]={'held_cross_entropy':float(objective(weights,held_X,held_prob)),
                   'held_brier':float(jnp.mean((jax.nn.sigmoid(held_X@weights)-held_prob)**2)),
                   'source_cross_entropy':float(objective(weights,source_X,source_targets))}
assert metrics['supervised']['held_cross_entropy']<metrics['pretrained']['held_cross_entropy']
assert metrics['teacher']['held_cross_entropy']<metrics['pretrained']['held_cross_entropy']
print(json.dumps(metrics,indent=2))
print('checkpoint source hash:',source_hash)

Code building blocks in this step

jnp.array(values, dtype=...) — Constructs an immutable device-backed JAX array from Python/NumPy values.
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.
jax.grad(loss_fn)(params, ...) — Transforms a scalar-output function into a function returning the gradient PyTree with the same structure as params.

The teacher is the known generating rule in this fixture, so matching it is an especially favorable distillation setting. Report source retention as well as target-task improvement.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

# Step 1 — 1. Pretrain a small model and retain its source data identity: The same stable cross-entropy accepts binary labels or teacher...
import hashlib
import json
import tempfile
from pathlib import Path
import numpy as np
import jax
import jax.numpy as jnp

def design(key,n):
    # Sample deterministic random values into `features` using an explicit PRNG key.
    features=jax.random.normal(key,(n,2))
    return jnp.concatenate([features,jnp.ones((n,1))],axis=1)
# Create or split explicit PRNG key(s) (`source_X`) for reproducible randomness.
source_X=design(jax.random.key(10),256)
source_targets=jax.nn.sigmoid(source_X@jnp.array([1.2,-.8,.2]))
# Create or split explicit PRNG key(s) (`adapt_X`) for reproducible randomness.
adapt_X=design(jax.random.key(11),48)
# Create or split explicit PRNG key(s) (`held_X`) for reproducible randomness.
held_X=design(jax.random.key(12),512)
teacher=jnp.array([.6,1.4,-.3])
soft_targets=jax.nn.sigmoid(adapt_X@teacher)
# Create or split explicit PRNG key(s) (`hard_targets`) for reproducible randomness.
hard_targets=jax.random.bernoulli(jax.random.key(13),soft_targets).astype(jnp.float32)
held_prob=jax.nn.sigmoid(held_X@teacher)
def objective(w,X,targets):
    logits=X@w
    return jnp.mean(jnp.logaddexp(0.,logits)-targets*logits)
def train(initial,X,targets,steps=300,rate=.15):
    def step(w,_):
        # Evaluate both scalar loss and parameter gradients in one pass (`(loss, grad)`).
        loss,grad=jax.value_and_grad(objective)(w,X,targets)
        return w-rate*grad,loss
    return jax.lax.scan(step,initial,None,length=steps)
# Allocate initialized array `(base, pretrain_history)` with the specified shape and dtype.
base,pretrain_history=train(jnp.zeros(3),source_X,source_targets)
assert pretrain_history[-1]<pretrain_history[0]-.1
# Convert `source_hash` to a host NumPy array for inspection or verification.
source_hash=hashlib.sha256(np.asarray(source_X).tobytes()+np.asarray(source_targets).tobytes()).hexdigest()

# Step 2 — 2. Save and reload before adapting: Both branches begin with identical pretrained parameters and a...
# Create an isolated temporary directory to run and inspect artifacts safely:
with tempfile.TemporaryDirectory() as directory:
    checkpoint=Path(directory)/'pretrained.npz'
    # Convert `` to a host NumPy array for inspection or verification.
    np.savez(checkpoint,weights=np.asarray(base))
    # Compute deterministic cryptographic digest `digest` for provenance verification.
    digest=hashlib.sha256(checkpoint.read_bytes()).hexdigest()
    # Compute deterministic cryptographic digest `manifest` for provenance verification.
    manifest={'model':'three-parameter logistic synthetic fixture','source_data_sha256':source_hash,
              'weights_sha256':digest,'source_seed':10,'pretrain_steps':300,'learning_rate':.15,
              'objective':'binary cross entropy with synthetic source probabilities',
              'jax':jax.__version__,'dtype':'float32','optimizer_state_included':False}
    manifest_path=Path(directory)/'manifest.json'
    # Write the serialized artifact payload to disk.
    manifest_path.write_text(json.dumps(manifest))
    recorded=json.loads(manifest_path.read_text())
    assert hashlib.sha256(checkpoint.read_bytes()).hexdigest()==recorded['weights_sha256']
    with np.load(checkpoint,allow_pickle=False) as saved:
        # Create device-backed JAX array `restored`.
        restored=jnp.asarray(saved['weights'])
# Execute `np.testing.assert_array_equal(restored,base)`
np.testing.assert_array_equal(restored,base)
frozen=np.array(restored,copy=True)
sft,sft_history=train(restored,adapt_X,hard_targets,steps=200)
distilled,distill_history=train(restored,adapt_X,soft_targets,steps=200)
# Execute `np.testing.assert_array_equal(restored,frozen)`
np.testing.assert_array_equal(restored,frozen)

# Step 3 — 3. Compare on a held-out criterion shared by both objectives: The teacher is the known generating rule in this fixture, so...
w_check=jnp.array([.2,-.3,.1])
# Convert `host_X` to a host NumPy array for inspection or verification.
host_X=np.asarray(adapt_X,dtype=np.float64)
# Convert `host_w` to a host NumPy array for inspection or verification.
host_w=np.asarray(w_check,dtype=np.float64)
host_p=1/(1+np.exp(-(host_X@host_w)))
for targets in (hard_targets,soft_targets):
    # Convert `expected` to a host NumPy array for inspection or verification.
    expected=host_X.T@(host_p-np.asarray(targets))/len(host_X)
    # Differentiate the objective to obtain gradients ``.
    np.testing.assert_allclose(jax.grad(objective)(w_check,adapt_X,targets),expected,rtol=1e-5,atol=1e-6)
metrics={}
for name,weights in [('pretrained',restored),('supervised',sft),('teacher',distilled)]:
    metrics[name]={'held_cross_entropy':float(objective(weights,held_X,held_prob)),
                   'held_brier':float(jnp.mean((jax.nn.sigmoid(held_X@weights)-held_prob)**2)),
                   'source_cross_entropy':float(objective(weights,source_X,source_targets))}
assert metrics['supervised']['held_cross_entropy']<metrics['pretrained']['held_cross_entropy']
assert metrics['teacher']['held_cross_entropy']<metrics['pretrained']['held_cross_entropy']
print(json.dumps(metrics,indent=2))
print('checkpoint source hash:',source_hash)

Run the example

python3 phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/code/main.py

About this command

This shortcut uses an activated macOS/Linux course environment in the repository folder. For Windows, use the explicit interpreter command in “Setup help” above. The script includes worked examples and reference solutions.

Inspect the complete assembled example

main.py · complete Python example

# Step 1 — 1. Pretrain a small model and retain its source data identity: The same stable cross-entropy accepts binary labels or teacher...
import hashlib
import json
import tempfile
from pathlib import Path
import numpy as np
import jax
import jax.numpy as jnp

def design(key,n):
    # Sample deterministic random values into `features` using an explicit PRNG key.
    features=jax.random.normal(key,(n,2))
    return jnp.concatenate([features,jnp.ones((n,1))],axis=1)
# Create or split explicit PRNG key(s) (`source_X`) for reproducible randomness.
source_X=design(jax.random.key(10),256)
source_targets=jax.nn.sigmoid(source_X@jnp.array([1.2,-.8,.2]))
# Create or split explicit PRNG key(s) (`adapt_X`) for reproducible randomness.
adapt_X=design(jax.random.key(11),48)
# Create or split explicit PRNG key(s) (`held_X`) for reproducible randomness.
held_X=design(jax.random.key(12),512)
teacher=jnp.array([.6,1.4,-.3])
soft_targets=jax.nn.sigmoid(adapt_X@teacher)
# Create or split explicit PRNG key(s) (`hard_targets`) for reproducible randomness.
hard_targets=jax.random.bernoulli(jax.random.key(13),soft_targets).astype(jnp.float32)
held_prob=jax.nn.sigmoid(held_X@teacher)
def objective(w,X,targets):
    logits=X@w
    return jnp.mean(jnp.logaddexp(0.,logits)-targets*logits)
def train(initial,X,targets,steps=300,rate=.15):
    def step(w,_):
        # Evaluate both scalar loss and parameter gradients in one pass (`(loss, grad)`).
        loss,grad=jax.value_and_grad(objective)(w,X,targets)
        return w-rate*grad,loss
    return jax.lax.scan(step,initial,None,length=steps)
# Allocate initialized array `(base, pretrain_history)` with the specified shape and dtype.
base,pretrain_history=train(jnp.zeros(3),source_X,source_targets)
assert pretrain_history[-1]<pretrain_history[0]-.1
# Convert `source_hash` to a host NumPy array for inspection or verification.
source_hash=hashlib.sha256(np.asarray(source_X).tobytes()+np.asarray(source_targets).tobytes()).hexdigest()

# Step 2 — 2. Save and reload before adapting: Both branches begin with identical pretrained parameters and a...
# Create an isolated temporary directory to run and inspect artifacts safely:
with tempfile.TemporaryDirectory() as directory:
    checkpoint=Path(directory)/'pretrained.npz'
    # Convert `` to a host NumPy array for inspection or verification.
    np.savez(checkpoint,weights=np.asarray(base))
    # Compute deterministic cryptographic digest `digest` for provenance verification.
    digest=hashlib.sha256(checkpoint.read_bytes()).hexdigest()
    # Compute deterministic cryptographic digest `manifest` for provenance verification.
    manifest={'model':'three-parameter logistic synthetic fixture','source_data_sha256':source_hash,
              'weights_sha256':digest,'source_seed':10,'pretrain_steps':300,'learning_rate':.15,
              'objective':'binary cross entropy with synthetic source probabilities',
              'jax':jax.__version__,'dtype':'float32','optimizer_state_included':False}
    manifest_path=Path(directory)/'manifest.json'
    # Write the serialized artifact payload to disk.
    manifest_path.write_text(json.dumps(manifest))
    recorded=json.loads(manifest_path.read_text())
    assert hashlib.sha256(checkpoint.read_bytes()).hexdigest()==recorded['weights_sha256']
    with np.load(checkpoint,allow_pickle=False) as saved:
        # Create device-backed JAX array `restored`.
        restored=jnp.asarray(saved['weights'])
# Execute `np.testing.assert_array_equal(restored,base)`
np.testing.assert_array_equal(restored,base)
frozen=np.array(restored,copy=True)
sft,sft_history=train(restored,adapt_X,hard_targets,steps=200)
distilled,distill_history=train(restored,adapt_X,soft_targets,steps=200)
# Execute `np.testing.assert_array_equal(restored,frozen)`
np.testing.assert_array_equal(restored,frozen)

# Step 3 — 3. Compare on a held-out criterion shared by both objectives: The teacher is the known generating rule in this fixture, so...
w_check=jnp.array([.2,-.3,.1])
# Convert `host_X` to a host NumPy array for inspection or verification.
host_X=np.asarray(adapt_X,dtype=np.float64)
# Convert `host_w` to a host NumPy array for inspection or verification.
host_w=np.asarray(w_check,dtype=np.float64)
host_p=1/(1+np.exp(-(host_X@host_w)))
for targets in (hard_targets,soft_targets):
    # Convert `expected` to a host NumPy array for inspection or verification.
    expected=host_X.T@(host_p-np.asarray(targets))/len(host_X)
    # Differentiate the objective to obtain gradients ``.
    np.testing.assert_allclose(jax.grad(objective)(w_check,adapt_X,targets),expected,rtol=1e-5,atol=1e-6)
metrics={}
for name,weights in [('pretrained',restored),('supervised',sft),('teacher',distilled)]:
    metrics[name]={'held_cross_entropy':float(objective(weights,held_X,held_prob)),
                   'held_brier':float(jnp.mean((jax.nn.sigmoid(held_X@weights)-held_prob)**2)),
                   'source_cross_entropy':float(objective(weights,source_X,source_targets))}
assert metrics['supervised']['held_cross_entropy']<metrics['pretrained']['held_cross_entropy']
assert metrics['teacher']['held_cross_entropy']<metrics['pretrained']['held_cross_entropy']
print(json.dumps(metrics,indent=2))
print('checkpoint source hash:',source_hash)

Expected result: Both adapted branches reduce the common held-out target cross-entropy relative to the reloaded source model. The printed table also reports probability error and source-task retention; inspect all three quantities.

Recorded reference execution

Reference output captured from executing the complete lesson script (code/main.py) and its assertions in the course CPU environment.

2026-10-20T10:50:24.147400+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

{
  "pretrained": {
    "held_cross_entropy": 0.9780182838439941,
    "held_brier": 0.17108705639839172,
    "source_cross_entropy": 0.5451400876045227
  },
  "supervised": {
    "held_cross_entropy": 0.5554870367050171,
    "held_brier": 0.012093935161828995,
    "source_cross_entropy": 0.7904530763626099
  },
  "teacher": {
    "held_cross_entropy": 0.5240296125411987,
    "held_brier": 0.00022982103109825402,
    "source_cross_entropy": 1.0046310424804688
  }
}
checkpoint source hash: f5fb50c38195b46766bee2edaa81a7ad28386523baa850c5294eb71b47f1f84c
{
  "pretrained": {
    "held_cross_entropy": 0.9780182838439941,
    "held_brier": 0.17108705639839172,
    "source_cross_entropy": 0.5451400876045227
  },
  "supervised": {
    "held_cross_entropy": 0.5554870367050171,
    "held_brier": 0.012093935161828995,
    "source_cross_entropy": 0.7904530763626099
  },
  "teacher": {
    "held_cross_entropy": 0.5240296125411987,
    "held_brier": 0.00022982103109825402,
    "source_cross_entropy": 1.0046310424804688
  }
}
checkpoint source hash: f5fb50c38195b46766bee2edaa81a7ad28386523baa850c5294eb71b47f1f84c
wrong-teacher held loss: 1.3151594400405884
source retention losses: [0.5451400876045227, 0.7904530763626099, 1.0046310424804688]
supervised source loss increase 0.24531298875808716 passes declared gate True
teacher source loss increase 0.45949095487594604 passes declared gate False
Candidate selected under the declared fixture budget: supervised
PASS: deployment-02

Download execution record · The notebook includes executed cells and figure outputs. Compare your run with these values; timings depend on the machine.

Check a stable extreme logit

Check a stable extreme logit · Python experiment

# Experiment — Check a stable extreme logit: logaddexp evaluates softplus stably; a naive log(1+exp(z)) can...
extreme=objective(jnp.array([1000.]),jnp.ones((1,1)),jnp.zeros(1))
assert jnp.isfinite(extreme) and jnp.allclose(extreme,1000.)

Expected: The wrong confident prediction has finite loss $1000$ .

logaddexp evaluates softplus stably; a naive log(1+exp(z)) can overflow.

Reveal teacher error

Reveal teacher error · Python experiment

# Experiment — Reveal teacher error: Distillation optimizes fidelity to its teacher, which is not...
wrong,_=train(restored,adapt_X,1-soft_targets,steps=200)
assert objective(wrong,held_X,held_prob)>objective(distilled,held_X,held_prob)
print("wrong-teacher held loss:",float(objective(wrong,held_X,held_prob)))

Expected: The deliberately wrong teacher produces worse target evaluation than the correct teacher.

Distillation optimizes fidelity to its teacher, which is not automatically fidelity to the desired task.

Hands-On Exercise · Modify & Verify Run in the declared lesson environment

Make it yours

Where to put your code: Append your solution at the bottom of phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/code/main.py (or inside the Make it yours cell of exercises/deployment-02.ipynb) and run python3 phases/15-deployment/02-adapt-a-pretrained-model-and-choose-a-post-training-objective/code/main.py to verify your prediction before opening the reference solution.

Adapt for zero steps and verify that predictions exactly match the reloaded checkpoint. Then run a short adaptation and show that the saved base is unchanged while the adapted parameters differ.

Show a writing guide and starter scaffold

How to write this exercise code: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

train(...) — Call train with your updated parameters or inputs from this lesson's workspace.
testing.assert_array_equal(...) — Call testing.assert_array_equal with your updated parameters or inputs from this lesson's workspace.
assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

2. Step-by-step implementation plan:

Execute np.testing.assert_array_equal(unchanged,restored)
Execute np.testing.assert_array_equal(restored,frozen)

Starter code scaffold · copy into your file & fill in the TODOs

# Exercise solution: Adapt for zero steps and verify that predictions exactly match the...
unchanged,_ = train(...)  # TODO: compute unchanged,_
# Execute `np.testing.assert_array_equal(unchanged,restored)`
np.testing.assert_array_equal(unchanged,restored)
short,_ = train(...)  # TODO: compute short,_
assert not jnp.array_equal(short,restored)  # TODO: complete assertion check
# Execute `np.testing.assert_array_equal(restored,frozen)`
np.testing.assert_array_equal(restored,frozen)

Reveal the reference solution

Reference solution

# Exercise solution: Adapt for zero steps and verify that predictions exactly match the...
unchanged,_=train(restored,adapt_X,hard_targets,steps=0)
# Execute `np.testing.assert_array_equal(unchanged,restored)`
np.testing.assert_array_equal(unchanged,restored)
short,_=train(restored,adapt_X,hard_targets,steps=10)
assert not jnp.array_equal(short,restored)
# Execute `np.testing.assert_array_equal(restored,frozen)`
np.testing.assert_array_equal(restored,frozen)

Keep source-task evidence

Transfer / diagnosis

Measure source cross-entropy before and after both adaptation objectives. Explain any forgetting using the different source and target probability rules.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

evidence(...) — Call evidence with your updated parameters or inputs from this lesson's workspace.
objective(...) — Call objective with your updated parameters or inputs from this lesson's workspace.
assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

2. Step-by-step implementation plan:

Set up the changed input or configuration parameters requested in the prompt.
Run the target function or transformation and assert that the output matches your prediction.

Starter code scaffold · copy into your file & fill in the TODOs

# Keep source-task evidence (Transfer / diagnosis): The two tasks prefer different parameters; the measurement...
retention = ...  # TODO: compute retention
assert retention[1]  # TODO: complete assertion check
print("source retention losses:",retention)

Show a hint

Report the actual values; target improvement is not a source-retention guarantee.

Reveal reference solution and reasoning

Keep source-task evidence · Python solution

# Keep source-task evidence (Transfer / diagnosis): The two tasks prefer different parameters; the measurement...
retention=[float(objective(w,source_X,source_targets)) for w in (restored,sft,distilled)]
assert retention[1]>retention[0] and retention[2]>retention[0]
print("source retention losses:",retention)

The two tasks prefer different parameters; the measurement exposes that tradeoff.

Check the loss without autodiff

Transfer / diagnosis

Compare the soft-target loss with a host float64 logaddexp expression on new weights.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).
jnp.allclose(actual, expected, rtol=..., atol=...) — Checks that two arrays match elementwise within floating-point tolerance.

2. Step-by-step implementation plan:

Perform matrix / vector contraction (@) to compute z.
Aggregate array values to compute expected.
Create device-backed JAX array .

Starter code scaffold · copy into your file & fill in the TODOs

# Check the loss without autodiff (Transfer / diagnosis): Independent algebra checks normalization and reduction as...
probe = np.array(...)  # TODO: compute probe
# Perform matrix / vector contraction (`@`) to compute `z`.
z = ...  # TODO: compute z
# Aggregate array values to compute `expected`.
expected = np.mean(...)  # TODO: compute expected
# Create device-backed JAX array ``.
np.testing.assert_allclose(objective(jnp.asarray(probe),adapt_X,soft_targets),expected,rtol = ...  # TODO: compute np.testing.assert_allclose(objective(jnp.asarray(probe),adapt_X,soft_targets),expected,rtol

Show a hint

Use the same logits, reduction and targets; this check does not call the JAX loss as its oracle.

Reveal reference solution and reasoning

Check the loss without autodiff · Python solution

# Check the loss without autodiff (Transfer / diagnosis): Independent algebra checks normalization and reduction as...
probe=np.array([-.4,.2,.3])
# Perform matrix / vector contraction (`@`) to compute `z`.
z=host_X@probe
# Aggregate array values to compute `expected`.
expected=np.mean(np.logaddexp(0,z)-np.asarray(soft_targets)*z)
# Create device-backed JAX array ``.
np.testing.assert_allclose(objective(jnp.asarray(probe),adapt_X,soft_targets),expected,rtol=1e-6)

Independent algebra checks normalization and reduction as well as stability.

Choose under a declared retention constraint

Transfer / diagnosis

Before running, choose a source-loss increase budget of $0.3$ nats for this fixture. Evaluate both adapted candidates against the same target and source criteria; retain the base if neither qualifies.

Show a writing guide and starter scaffold

How to write this practice solution: step-by-step recipe & starter scaffold

1. Key functions & syntax to use:

constraint(...) — Call constraint with your updated parameters or inputs from this lesson's workspace.
in(...) — Call in with your updated parameters or inputs from this lesson's workspace.
assert condition — Verify that the observed output shape, status, or numerical value satisfies the contract.

2. Step-by-step implementation plan:

Set up the changed input or configuration parameters requested in the prompt.
Run the target function or transformation and assert that the output matches your prediction.

Starter code scaffold · copy into your file & fill in the TODOs

# Choose under a declared retention constraint (Transfer / diagnosis): The gate separates optimization from acceptance.
retention_budget = ...  # TODO: compute retention_budget
eligible = ...  # TODO: compute eligible
for name in ('supervised', 'teacher'):
    delta = ...  # TODO: compute delta
    improved = ...  # TODO: compute improved
    accepted_candidate = ...  # TODO: compute accepted_candidate
    print(name, 'source loss increase', delta, 'passes declared gate', accepted_candidate)
    if accepted_candidate: eligible.append(name)
selected = min(...)  # TODO: compute selected
assert selected  # TODO: complete assertion check
assert all(metrics[n]['source_cross_entropy']-metrics['pretrained']['source_cross_entropy']  # TODO: complete assertion check
print('Candidate selected under the declared fixture budget:', selected)

Show a hint

Compare each source loss with the base source loss. Do not compare the two training objectives with each other.

Reveal reference solution and reasoning

Choose under a declared retention constraint · Python solution

# Choose under a declared retention constraint (Transfer / diagnosis): The gate separates optimization from acceptance.
retention_budget = .3
eligible = []
for name in ('supervised', 'teacher'):
    delta = metrics[name]['source_cross_entropy'] - metrics['pretrained']['source_cross_entropy']
    improved = metrics[name]['held_cross_entropy'] < metrics['pretrained']['held_cross_entropy']
    accepted_candidate = improved and delta <= retention_budget
    print(name, 'source loss increase', delta, 'passes declared gate', accepted_candidate)
    if accepted_candidate: eligible.append(name)
selected = min(eligible, key=lambda n: metrics[n]['held_cross_entropy']) if eligible else 'pretrained'
assert selected == 'pretrained' or selected in eligible
assert all(metrics[n]['source_cross_entropy']-metrics['pretrained']['source_cross_entropy'] <= retention_budget for n in eligible)
print('Candidate selected under the declared fixture budget:', selected)

The gate separates optimization from acceptance. Its conclusion applies to these seeded data and this declared budget; selecting a candidate is not evidence of generalization to a natural dataset.

Diagnose the result

If target performance improves but source loss rises, inspect task conflict before blaming checkpoint corruption. If the initial adapted output differs before any update, compare checkpoint hash, preprocessing and precision. If a distillation result looks implausibly strong, verify that teacher labels and evaluation labels were not generated from the same held-out examples.

Carry forward

A post-training comparison begins with an identifiable trained artifact.
Choose an objective according to the supervision actually available.
Compare target quality, source retention and provenance before exporting a new version.

Keep your evidence

Download evidence template

Source/checkpoint hashes, independent loss and gradient, disjoint evaluation results and teacher-failure diagnosis. Keep the environment, observed outputs and your explanation of the figure.

Save your modified code, its output, and an explanation of your prediction. Checkpoint answers and exercise evidence are tracked separately.

Saved in this browser. Export or back up your portfolio anytime in My learning.

Save work for this lesson

Verified on local CPU. Stage to a Cloud TPU VM via Cloud TPU setup to record accelerator execution receipts (TODO: add automated Cloud TPU CI execution receipts).

Adapt a pretrained model and choose a post-training objective

What you will be able to do

The problem

The idea

Choose the target before choosing the adaptation method

Define the tasks before choosing an objective

Understand hard labels and soft targets

Make pretrained mean previously trained

Choose a fair comparison

Locate preference and reinforcement objectives

Carry the contract to a real pretrained model

Turn task tradeoffs into a release decision

Two adaptation objectives evaluated on the same held-out target

Read the figure

Connect it to the computation

1 Pretrain a small model and retain its source data identity

2 Save and reload before adapting

3 Compare on a held-out criterion shared by both objectives

Run the example

Recorded reference execution

Check a stable extreme logit

Reveal teacher error

Make it yours

Keep source-task evidence

Check the loss without autodiff

Choose under a declared retention constraint

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references