Training Systems · Local CPULaptop CPULesson 2 of 3 in Phase 17

Phase 17 / Self-supervised pretraining: masked and contrastive learning

Masked image modeling: reconstruct missing patches

Revision v0.7.0 · 2026-10-20

Lesson · about 105 minutes · CPU practice

Preserve patch order and pixel order. Keep hidden pixels out of the encoder.

Phase study guide: preparation, milestones and project evidence

Before you start

Before this: Masked language modeling: predict hidden tokens
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/17-pretraining/02-mim/code/main.py

PowerShell command · repository folder

.\.venv\Scripts\python.exe phases/17-pretraining/02-mim/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 pretraining-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: Hidden-patch training MSE falls below 0.002 and held-amplitude MSE stays below 0.02.

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

Phase 17.02 Stage 1 of 5 · Understand

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

What you will be able to do

Preserve patch order and pixel order
Keep hidden pixels out of the encoder
Verify the changed-condition exercise and explain the limits of this fixture.

The problem

If part of an image is hidden, can the model infer it from visible patches? The important question is whether hidden pixels are truly absent from the encoder, not whether the reconstruction looks attractive.

The idea

Masked image modeling learns from visible image content to predict missing content. To understand the objective, distinguish observed patches, hidden targets and predicted values in the original image coordinates.

Keep hidden patches in their spatial positions

The fixture splits a four-by-four image into four two-by-two patches in row-major order. Patches $0$ and $3$ are visible; $1$ and $2$ are hidden. Show hidden regions with hatching so they cannot be mistaken for observed black pixels.

The implemented reconstructor is linear: eight visible pixel values plus a constant feature map through a nine-by-sixteen weight matrix. It is a deliberately small mechanism example, not a masked autoencoder Transformer.

Compare target and predicted hidden patches at the same spatial coordinates and on a shared intensity scale. A separate error map reveals remaining mistakes that a small average loss can hide. Changing a hidden target without changing visible input must not change the model's forward input.

The sixteen cells preserve the four-by-four image layout. Green patches 0 and 3 provide visible values. Hatched patches 1 and 2 are hidden targets for reconstruction loss. Hatching means missing input, not observed black pixels. This maps the actual mask without inventing a reconstructed image or error measurement. — Architecture and dataflow mechanism diagram.

Before coding: what information can the model use?

Imagine covering the top-right and bottom-left quarters of a small image. The model is allowed to see the other two quarters and must infer the covered pixels. Keep two separate paths on your sketch: visible patches go into features, while the untouched image supplies targets. A zero loss is suspicious if the target image is also an encoder input.

This dataset varies a single brightness amplitude on top of a fixed spatial pattern. Visible pixels reveal that amplitude, so a linear map can succeed. That makes the example suitable for auditing masking and patch order, but it cannot teach natural-image representation quality on its own.

Preserve patch order and pixel order

Images use batch, height, width, channel axes. A $4\times4$ single-channel image becomes four $2\times2$ patches with four scalars each. Reshape alone does not generally produce row-major patch order; transpose the grid and within-patch axes deliberately. The first patch is checked against an independently sliced image region.

Keep hidden pixels out of the encoder

Gather visible patches before constructing features. The clean full image belongs only to the reconstruction target. In a masked autoencoder, a visible-only encoder can reduce encoder work while a decoder uses position information and mask tokens to restore the complete patch ordering. Our linear decoder has an explicit output coordinate for each patch/pixel, so it has a fixed positional contract without a Transformer.

Work through a patch rather than memorizing reshape

Number a single-channel image row by row from $0$ through $15$ . Its top-left patch is $[0,1,4,5]$ , not $[0,1,2,3]$ . The latter is a horizontal strip caused by flattening without the grid-axis transpose. For general image shape $(B,H,W,C)$ and patch width $p$ , the output shape is $(B,HW/p^2,p^2C)$ . The channel coordinate remains inside each pixel.

In our fixture, two visible patches each contribute four pixels; appending a bias gives nine features. The decoder has sixteen outputs because it predicts every patch coordinate. Its visible outputs are unconstrained by the hidden-only objective. A whole-image picture should therefore show a composite that preserves visible input pixels and inserts predictions only where hidden.

Understand which uncertainty the objective can solve

If two examples have exactly the same visible pixels but different hidden targets, no deterministic model can reconstruct both perfectly. For a scalar target equally likely to be $0$ or $2$ , the mean-squared-error optimum is $1$ , with average squared error $1$ . The result looks like neither original target. More capacity cannot recover information absent from the inputs.

This explains why blurry reconstructions need not mean an implementation bug. First check information flow, patch layout and the irreducible ambiguity of the data; then investigate capacity or optimization. A useful downstream representation must be tested separately from pixel reconstruction.

Reduce over the hidden reconstruction targets

First average squared pixel error within each patch, then sum the hidden-patch errors and divide by the hidden count. This definition gives equal weight to every hidden pixel because all patches have equal size. Errors on visible outputs are excluded. Mixing visible and hidden pixels can flatter the score when copying visible inputs is easy.

L=\frac{\sum_{b,k}m_{bk}\,\frac1P\sum_{j=1}^{P}(\widehat x_{bkj}-x_{bkj})^2}{\sum_{b,k}m_{bk}}

Interpret normalization and reconstruction

Raw-pixel reconstruction and per-patch normalized reconstruction are different objectives. Normalization can remove a patch’s brightness/contrast information and requires a clear inverse or evaluation convention. This lesson uses raw synthetic values and records MSE in squared input units. Do not compare that number directly with a differently normalized experiment.

Evaluate beyond the training mask

The reference checks intermediate amplitudes that were not training examples. That is a small interpolation test, not semantic recognition. For transfer, freeze the encoder and train a probe on separate labels, or fine-tune under a held-out protocol. Compare mask ratios with fixed compute/data budgets and keep all-hidden or all-visible policies explicit.

Make a transfer test that can falsify your explanation

Our held-out amplitudes test interpolation under the same fixed spatial pattern. They do not test new mask layouts. Changing the visible positions while keeping the same feature coordinate meanings silently changes the model contract. A variable-mask encoder needs positions and a visible-token gathering/restoration scheme.

Before claiming a transferable image encoder, test new source images, compare a frozen probe and fine-tuning with a matched baseline, and inspect class-specific failures. In this lesson, keep a simpler evidence packet: a numbered patch oracle, a hidden-pixel intervention, an ambiguity example, a composite reconstruction, and a written explanation of which of these tests would catch leakage.

Masked image modeling: reconstruct missing patches — recorded experiment

Predict before looking: Will every predicted pixel match its target exactly? Which output patches were never trained?

Read the figure

The line shows hidden-patch training MSE before each update, falling from about $0.554$ to $0.00169$ . The vertical units are squared raw pixel units. The line excludes visible-patch reconstruction errors; it cannot be read as full-image MSE.

The second panel inspects the first image only. Its first two rows are clean hidden-patch targets; the last two are final predictions in the same within-patch order. Compare row one with row three and row two with row four, rather than looking for a diagonal. Residual brightness differences explain why the training MSE is small but nonzero. Visible outputs are omitted because this loss never trains them; displaying them as a full reconstructed image would be misleading.

Connect it to the computation

Each image contains an amplitude and a fixed spatial pattern. Visible patches reveal the amplitude, so a linear decoder can infer hidden pixels on this deliberately learnable task. Separate held-amplitude MSE is approximately 0.00118. Natural-image semantics and a visible-token Transformer remain different experiments.

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/17-pretraining/02-mim/code/main.py (or downloaded exercises/pretraining-02.py / exercises/pretraining-02.ipynb) and execute from the repository root with python3 phases/17-pretraining/02-mim/code/main.py.

1 Define patch layout and masked error

Create main.py in your activated course environment. Paste this block, then run python main.py; function definitions alone print nothing.

main.py · append build step 1

# Step 1 — 1. Define patch layout and masked error: Trace the transpose against the numbered top-left patch.
import jax
import jax.numpy as jnp
import numpy as np

def patchify(images, patch=2):
    b, h, w, c = images.shape
    # Guard input contract (`h % patch or w % patch`) and fail fast if violated.
    if h % patch or w % patch:
        raise ValueError('image dimensions must be divisible by patch size')
    return (
        images.reshape(b, h // patch, patch, w // patch, patch, c)
        .transpose(0, 1, 3, 2, 4, 5)
        .reshape(b, (h // patch) * (w // patch), patch * patch * c)
    )

def masked_mse(prediction, target, hidden):
    # Reduce along axis=-1 to compute `per_patch`.
    per_patch = jnp.mean((prediction - target) ** 2, axis=-1)
    return jnp.sum(jnp.where(hidden, per_patch, 0.0)) / jnp.sum(hidden)

def validate_mask(selected, shape):
    # Convert `a` to a host NumPy array for inspection or verification.
    a = np.asarray(selected)
    # Guard input contract (`a.shape != shape or a.dtype != np.bool_ or (not a.any())`) and fail fast if violated.
    if a.shape != shape or a.dtype != np.bool_ or not a.any():
        raise ValueError('a nonempty Boolean mask of the target shape is required')

Code building blocks in this step

array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).

Trace the transpose against the numbered top-left patch. masked_mse gives equal weight to each hidden pixel only because every patch has equal size.

2 Separate visible features from full targets

Append this block to main.py and run python main.py again. Keep the earlier blocks above it.

main.py · append build step 2

# Four patches share an amplitude plus a fixed spatial offset; no real photographs.
amplitude = jnp.array([0.1, 0.3, 0.6, 0.8])
# Construct and reshape `base_image` into the target tensor dimensions.
base_image = jnp.arange(16, dtype=jnp.float32).reshape(4, 4, 1) / 32
images = amplitude[:, None, None, None] + base_image[None]
patches = patchify(images)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(patches[0, 0], np.asarray(images[0, :2, :2, :]).reshape(-1))
hidden = jnp.array([[False, True, True, False]] * 4)
validate_mask(hidden, patches.shape[:2])
# The encoder receives only visible patches; original hidden pixels are targets only.
visible = patches[:, [0, 3], :]
# Construct and reshape `features` into the target tensor dimensions.
features = jnp.concatenate([visible.reshape(4, -1), jnp.ones((4, 1))], axis=1)
w = jnp.zeros((9, 16))
history = []
# Construct and reshape `loss` into the target tensor dimensions.
loss = lambda w: masked_mse((features @ w).reshape(4, 4, 4), patches, hidden)
# Differentiate the objective to obtain `step` via automatic differentiation.
step = jax.jit(jax.value_and_grad(loss))

Code building blocks in this step

jnp.arange(n, dtype=...) — Creates a 1-D JAX array of evenly spaced values [0, 1, ..., n-1] on the target device.
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.
array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).

Only patches $0$ and $3$ enter features. The bias permits a spatial offset; hidden pixels remain solely in patches, the target array.

3 Fit the decoder and audit the reconstruction

Append this block to main.py and run python main.py again. Keep the earlier blocks above it.

main.py · append build step 3

# Step 3 — 3. Fit the decoder and audit the reconstruction: The host loop checks the reduction independently.
for _ in range(100):
    value, g = step(w)
    history.append(float(value))
    w = w - 0.15 * g

# Construct and reshape `prediction` into the target tensor dimensions.
prediction = (features @ w).reshape(4, 4, 4)
assert history[-1] < history[0] * 0.02
# A scalar host loop is independent of the vectorized reduction.
expected = np.mean(
    [
        float(np.mean((np.asarray(prediction[b, k]) - np.asarray(patches[b, k])) ** 2))
        for b in range(4)
        for k in [1, 2]
    ]
)
np.testing.assert_allclose(loss(w), expected, rtol=1e-5)
visible_changed = prediction.at[:, [0, 3], :].set(999.0)
# Execute `np.testing.assert_allclose(`
np.testing.assert_allclose(
    masked_mse(visible_changed, patches, hidden), loss(w), atol=1e-7
)
held_images = jnp.array([0.2, 0.5])[:, None, None, None] + base_image[None]
held_patches = patchify(held_images)
# Construct and reshape `held_features` into the target tensor dimensions.
held_features = jnp.concatenate(
    [held_patches[:, [0, 3], :].reshape(2, -1), jnp.ones((2, 1))], axis=1
)
# Construct and reshape `held_loss` into the target tensor dimensions.
held_loss = float(
    masked_mse((held_features @ w).reshape(2, 4, 4), held_patches, hidden[:2])
)
assert held_loss < 0.02
print(
    'Masked pixel MSE initial/final/held amplitudes:',
    history[0],
    history[-1],
    held_loss,
)

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.
array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).
x.sum(axis=...) / x.mean(axis=..., keepdims=...) — Reduces values along the named axis (the axis you name is collapsed unless keepdims=True).

The host loop checks the reduction independently. Corrupting visible outputs checks the loss boundary; held amplitudes check a narrow interpolation claim.

View the cumulative build in main.py

These blocks append in order to the same file.

main.py · cumulative Python build

# Step 1 — 1. Define patch layout and masked error: Trace the transpose against the numbered top-left patch.
import jax
import jax.numpy as jnp
import numpy as np

def patchify(images, patch=2):
    b, h, w, c = images.shape
    # Guard input contract (`h % patch or w % patch`) and fail fast if violated.
    if h % patch or w % patch:
        raise ValueError('image dimensions must be divisible by patch size')
    return (
        images.reshape(b, h // patch, patch, w // patch, patch, c)
        .transpose(0, 1, 3, 2, 4, 5)
        .reshape(b, (h // patch) * (w // patch), patch * patch * c)
    )

def masked_mse(prediction, target, hidden):
    # Reduce along axis=-1 to compute `per_patch`.
    per_patch = jnp.mean((prediction - target) ** 2, axis=-1)
    return jnp.sum(jnp.where(hidden, per_patch, 0.0)) / jnp.sum(hidden)

def validate_mask(selected, shape):
    # Convert `a` to a host NumPy array for inspection or verification.
    a = np.asarray(selected)
    # Guard input contract (`a.shape != shape or a.dtype != np.bool_ or (not a.any())`) and fail fast if violated.
    if a.shape != shape or a.dtype != np.bool_ or not a.any():
        raise ValueError('a nonempty Boolean mask of the target shape is required')

# Four patches share an amplitude plus a fixed spatial offset; no real photographs.
amplitude = jnp.array([0.1, 0.3, 0.6, 0.8])
# Construct and reshape `base_image` into the target tensor dimensions.
base_image = jnp.arange(16, dtype=jnp.float32).reshape(4, 4, 1) / 32
images = amplitude[:, None, None, None] + base_image[None]
patches = patchify(images)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(patches[0, 0], np.asarray(images[0, :2, :2, :]).reshape(-1))
hidden = jnp.array([[False, True, True, False]] * 4)
validate_mask(hidden, patches.shape[:2])
# The encoder receives only visible patches; original hidden pixels are targets only.
visible = patches[:, [0, 3], :]
# Construct and reshape `features` into the target tensor dimensions.
features = jnp.concatenate([visible.reshape(4, -1), jnp.ones((4, 1))], axis=1)
w = jnp.zeros((9, 16))
history = []
# Construct and reshape `loss` into the target tensor dimensions.
loss = lambda w: masked_mse((features @ w).reshape(4, 4, 4), patches, hidden)
# Differentiate the objective to obtain `step` via automatic differentiation.
step = jax.jit(jax.value_and_grad(loss))

# Step 3 — 3. Fit the decoder and audit the reconstruction: The host loop checks the reduction independently.
for _ in range(100):
    value, g = step(w)
    history.append(float(value))
    w = w - 0.15 * g

# Construct and reshape `prediction` into the target tensor dimensions.
prediction = (features @ w).reshape(4, 4, 4)
assert history[-1] < history[0] * 0.02
# A scalar host loop is independent of the vectorized reduction.
expected = np.mean(
    [
        float(np.mean((np.asarray(prediction[b, k]) - np.asarray(patches[b, k])) ** 2))
        for b in range(4)
        for k in [1, 2]
    ]
)
np.testing.assert_allclose(loss(w), expected, rtol=1e-5)
visible_changed = prediction.at[:, [0, 3], :].set(999.0)
# Execute `np.testing.assert_allclose(`
np.testing.assert_allclose(
    masked_mse(visible_changed, patches, hidden), loss(w), atol=1e-7
)
held_images = jnp.array([0.2, 0.5])[:, None, None, None] + base_image[None]
held_patches = patchify(held_images)
# Construct and reshape `held_features` into the target tensor dimensions.
held_features = jnp.concatenate(
    [held_patches[:, [0, 3], :].reshape(2, -1), jnp.ones((2, 1))], axis=1
)
# Construct and reshape `held_loss` into the target tensor dimensions.
held_loss = float(
    masked_mse((held_features @ w).reshape(2, 4, 4), held_patches, hidden[:2])
)
assert held_loss < 0.02
print(
    'Masked pixel MSE initial/final/held amplitudes:',
    history[0],
    history[-1],
    held_loss,
)

Run the example

python3 phases/17-pretraining/02-mim/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. Define patch layout and masked error: Trace the transpose against the numbered top-left patch.
import jax
import jax.numpy as jnp
import numpy as np

def patchify(images, patch=2):
    b, h, w, c = images.shape
    # Guard input contract (`h % patch or w % patch`) and fail fast if violated.
    if h % patch or w % patch:
        raise ValueError('image dimensions must be divisible by patch size')
    return (
        images.reshape(b, h // patch, patch, w // patch, patch, c)
        .transpose(0, 1, 3, 2, 4, 5)
        .reshape(b, (h // patch) * (w // patch), patch * patch * c)
    )

def masked_mse(prediction, target, hidden):
    # Reduce along axis=-1 to compute `per_patch`.
    per_patch = jnp.mean((prediction - target) ** 2, axis=-1)
    return jnp.sum(jnp.where(hidden, per_patch, 0.0)) / jnp.sum(hidden)

def validate_mask(selected, shape):
    # Convert `a` to a host NumPy array for inspection or verification.
    a = np.asarray(selected)
    # Guard input contract (`a.shape != shape or a.dtype != np.bool_ or (not a.any())`) and fail fast if violated.
    if a.shape != shape or a.dtype != np.bool_ or not a.any():
        raise ValueError('a nonempty Boolean mask of the target shape is required')

# Four patches share an amplitude plus a fixed spatial offset; no real photographs.
amplitude = jnp.array([0.1, 0.3, 0.6, 0.8])
# Construct and reshape `base_image` into the target tensor dimensions.
base_image = jnp.arange(16, dtype=jnp.float32).reshape(4, 4, 1) / 32
images = amplitude[:, None, None, None] + base_image[None]
patches = patchify(images)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(patches[0, 0], np.asarray(images[0, :2, :2, :]).reshape(-1))
hidden = jnp.array([[False, True, True, False]] * 4)
validate_mask(hidden, patches.shape[:2])
# The encoder receives only visible patches; original hidden pixels are targets only.
visible = patches[:, [0, 3], :]
# Construct and reshape `features` into the target tensor dimensions.
features = jnp.concatenate([visible.reshape(4, -1), jnp.ones((4, 1))], axis=1)
w = jnp.zeros((9, 16))
history = []
# Construct and reshape `loss` into the target tensor dimensions.
loss = lambda w: masked_mse((features @ w).reshape(4, 4, 4), patches, hidden)
# Differentiate the objective to obtain `step` via automatic differentiation.
step = jax.jit(jax.value_and_grad(loss))

# Step 3 — 3. Fit the decoder and audit the reconstruction: The host loop checks the reduction independently.
for _ in range(100):
    value, g = step(w)
    history.append(float(value))
    w = w - 0.15 * g

# Construct and reshape `prediction` into the target tensor dimensions.
prediction = (features @ w).reshape(4, 4, 4)
assert history[-1] < history[0] * 0.02
# A scalar host loop is independent of the vectorized reduction.
expected = np.mean(
    [
        float(np.mean((np.asarray(prediction[b, k]) - np.asarray(patches[b, k])) ** 2))
        for b in range(4)
        for k in [1, 2]
    ]
)
np.testing.assert_allclose(loss(w), expected, rtol=1e-5)
visible_changed = prediction.at[:, [0, 3], :].set(999.0)
# Execute `np.testing.assert_allclose(`
np.testing.assert_allclose(
    masked_mse(visible_changed, patches, hidden), loss(w), atol=1e-7
)
held_images = jnp.array([0.2, 0.5])[:, None, None, None] + base_image[None]
held_patches = patchify(held_images)
# Construct and reshape `held_features` into the target tensor dimensions.
held_features = jnp.concatenate(
    [held_patches[:, [0, 3], :].reshape(2, -1), jnp.ones((2, 1))], axis=1
)
# Construct and reshape `held_loss` into the target tensor dimensions.
held_loss = float(
    masked_mse((held_features @ w).reshape(2, 4, 4), held_patches, hidden[:2])
)
assert held_loss < 0.02
print(
    'Masked pixel MSE initial/final/held amplitudes:',
    history[0],
    history[-1],
    held_loss,
)

Expected result: Hidden-patch training MSE falls below 0.002 and held-amplitude MSE stays below 0.02.

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:52:02.405713+00:00 · Python 3.14.3 · JAX 0.9.2 · 1 CPU device

Masked pixel MSE initial/final/held amplitudes: 0.5538085699081421 0.0016928859986364841 0.0011766081443056464
Masked pixel MSE initial/final/held amplitudes: 0.5538085699081421 0.0016928859986364841 0.0011766081443056464
Hidden MSE unchanged; full-image MSE is large.
Ambiguous target MSE at predictions 0, 1, 2: 2, 1, 2
Bottom-right patch order verified.
The zero-loss shortcut receives hidden pixels; the real features contain only eight visible pixels and a bias.
All rectangular RGB patches agree; hidden-pixel intervention preserves visible features.
PASS: pretraining-02

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

Make visible-output errors arbitrarily large

Make visible-output errors arbitrarily large · Python experiment

# Experiment — Make visible-output errors arbitrarily large: The disagreement identifies the reduction contract.
bad_visible = prediction.at[:, [0, 3], :].set(-1000.0)
# Execute `np.testing.assert_allclose(`
np.testing.assert_allclose(
    masked_mse(bad_visible, patches, hidden),
    masked_mse(prediction, patches, hidden),
)
assert float(jnp.mean((bad_visible - patches) ** 2)) > 1000
print('Hidden MSE unchanged; full-image MSE is large.')

Expected: Hidden MSE stays unchanged while full-image MSE grows dramatically.

The disagreement identifies the reduction contract. Neither metric is universally correct; label which region and preprocessing it measures.

Make hidden information impossible to recover

Make hidden information impossible to recover · Python experiment

# Experiment — Make hidden information impossible to recover: The experiment isolates missing information.
ambiguous_targets = jnp.array([[[0.0]], [[2.0]]])
ambiguous_mask = jnp.ones((2, 1), dtype=bool)

def ambiguity_loss(value):
    return masked_mse(
        jnp.full_like(ambiguous_targets, value), ambiguous_targets, ambiguous_mask
    )

# Execute `np.testing.assert_allclose(`
np.testing.assert_allclose(
    [ambiguity_loss(0.0), ambiguity_loss(1.0), ambiguity_loss(2.0)], [2.0, 1.0, 2.0]
)
# Differentiate the objective to obtain gradients ``.
np.testing.assert_allclose(jax.grad(ambiguity_loss)(1.0), 0.0, atol=1e-7)
print('Ambiguous target MSE at predictions 0, 1, 2: 2, 1, 2')

Expected: Predicting the mean is optimal but still leaves MSE of one.

The experiment isolates missing information. A loss plateau can be caused by ambiguity even when gradients and optimization are correct.

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/17-pretraining/02-mim/code/main.py (or inside the Make it yours cell of exercises/pretraining-02.ipynb) and run python3 phases/17-pretraining/02-mim/code/main.py to verify your prediction before opening the reference solution.

Verify the bottom-right patch against a direct slice and explain why a whole-array reshape can hide a spatial permutation.

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:

array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).

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

# Exercise solution: Verify the bottom-right patch against a direct slice and explain why a...
np.testing.assert_array_equal(
    patches[2, 3], np.asarray(images[2, 2:, 2:, :]).reshape(-1)
)
print('Bottom-right patch order verified.')

Reveal the reference solution

Reference solution

# Exercise solution: Verify the bottom-right patch against a direct slice and explain why a...
np.testing.assert_array_equal(
    patches[2, 3], np.asarray(images[2, 2:, 2:, :]).reshape(-1)
)
print('Bottom-right patch order verified.')

Expose the leakage shortcut

Transfer

Compare a decoder that returns the clean target with the visible-only model. Explain why zero reconstruction loss is not sufficient evidence.

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:

shortcut(...) — Call shortcut with your updated parameters or inputs from this lesson's workspace.
masked_mse(...) — Call masked_mse 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:

Check tensor shape invariant: features.shape[1] == 9

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

# Expose the leakage shortcut (Transfer): An information-flow check is necessary because the leaked...
shortcut = ...  # TODO: compute shortcut
assert float(masked_mse(shortcut, patches, hidden))  # TODO: complete assertion check
# Check tensor shape invariant: `features.shape[1] == 9`
assert features.shape[1]  # TODO: complete assertion check
print('The zero-loss shortcut receives hidden pixels; the real features contain only eight visible pixels and a bias.')

Show a hint

Measure the shortcut, then inspect which function inputs contain hidden pixels.

Reveal reference solution and reasoning

Expose the leakage shortcut · Python solution

# Expose the leakage shortcut (Transfer): An information-flow check is necessary because the leaked...
shortcut = patches
assert float(masked_mse(shortcut, patches, hidden)) == 0.0
# Check tensor shape invariant: `features.shape[1] == 9`
assert features.shape[1] == 9
print('The zero-loss shortcut receives hidden pixels; the real features contain only eight visible pixels and a bias.')

An information-flow check is necessary because the leaked solution can outperform the honest model on the objective itself.

Audit every patch of a rectangular color image

Challenge

Construct two 4-by-6 RGB images with numbered entries. Compare patchify with an independent nested slicing oracle for every patch, then change hidden pixels and prove visible features stay fixed.

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:

array.reshape(new_shape) — Reorganizes tensor axes without changing the total element count (array.size).

2. Step-by-step implementation plan:

Construct and reshape numbered into the target tensor dimensions.
Construct and reshape oracle into the target tensor dimensions.
Create device-backed JAX array .
Execute np.testing.assert_array_equal(changed_patches[:, [0, 3], :],

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

# Audit every patch of a rectangular color image (Challenge): The rectangular multichannel case catches axis errors that a...
# Construct and reshape `numbered` into the target tensor dimensions.
numbered = np.arange(...)  # TODO: compute numbered
# Construct and reshape `oracle` into the target tensor dimensions.
oracle = np.stack(...)  # TODO: compute oracle
    [
        np.stack(
            [
                image[row : row + 2, col : col + 2, :].reshape(-1)
                for row in range(0, 4, 2)
                for col in range(0, 6, 2)
            ]
        )
        for image in numbered
    ]
)
# Create device-backed JAX array ``.
np.testing.assert_array_equal(patchify(jnp.asarray(numbered)), oracle)
changed_images = ...  # TODO: compute changed_images
changed_patches = patchify(...)  # TODO: compute changed_patches
# Execute `np.testing.assert_array_equal(changed_patches[:, [0, 3], :],`
np.testing.assert_array_equal(changed_patches[:, [0, 3], :], visible)
assert not np.array_equal(  # TODO: complete assertion check
    np.asarray(changed_patches[:, [1, 2], :]), np.asarray(patches[:, [1, 2], :])
)
print('All rectangular RGB patches agree; hidden-pixel intervention preserves visible features.')

Show a hint

Use explicit image slices inside the oracle; do not copy the reshape/transpose implementation.

Reveal reference solution and reasoning

Audit every patch of a rectangular color image · Python solution

# Audit every patch of a rectangular color image (Challenge): The rectangular multichannel case catches axis errors that a...
# Construct and reshape `numbered` into the target tensor dimensions.
numbered = np.arange(2 * 4 * 6 * 3, dtype=np.float32).reshape(2, 4, 6, 3)
# Construct and reshape `oracle` into the target tensor dimensions.
oracle = np.stack(
    [
        np.stack(
            [
                image[row : row + 2, col : col + 2, :].reshape(-1)
                for row in range(0, 4, 2)
                for col in range(0, 6, 2)
            ]
        )
        for image in numbered
    ]
)
# Create device-backed JAX array ``.
np.testing.assert_array_equal(patchify(jnp.asarray(numbered)), oracle)
changed_images = images.at[:, :2, 2:, :].add(17.0).at[:, 2:, :2, :].add(-9.0)
changed_patches = patchify(changed_images)
# Execute `np.testing.assert_array_equal(changed_patches[:, [0, 3], :],`
np.testing.assert_array_equal(changed_patches[:, [0, 3], :], visible)
assert not np.array_equal(
    np.asarray(changed_patches[:, [1, 2], :]), np.asarray(patches[:, [1, 2], :])
)
print('All rectangular RGB patches agree; hidden-pixel intervention preserves visible features.')

The rectangular multichannel case catches axis errors that a single square grayscale patch can miss. The intervention independently checks that hidden targets cannot enter this encoder.

Diagnose the result

If reconstruction is excellent before learning, inspect target leakage. If patches look shuffled, verify the patch transpose and decoder restoration order. If metrics differ, compare the mask count and normalization.

Carry forward

Verify patch extraction against an explicit spatial slice: converting $(B,H,W,C)$ into $(B,HW/p^2,p^2C)$ requires transposing the grid and within-patch axes before flattening.
Keep hidden patches out of the encoder feature tensor and reduce reconstruction MSE over the hidden mask only; changing hidden target values must never alter the encoder forward input.

Keep your evidence

Download evidence template

Keep patch-order slices, visible-only feature shapes, hidden-patch counts, independent MSE and held-amplitude reconstruction results.

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).

Primary references

Masked Autoencoders ↗

Masked image modeling: reconstruct missing patches

What you will be able to do

The problem

The idea

Keep hidden patches in their spatial positions

Visible patches and hidden targets

Before coding: what information can the model use?

Preserve patch order and pixel order

Keep hidden pixels out of the encoder

Work through a patch rather than memorizing reshape

Understand which uncertainty the objective can solve

Reduce over the hidden reconstruction targets

Interpret normalization and reconstruction

Evaluate beyond the training mask

Make a transfer test that can falsify your explanation

Masked image modeling: reconstruct missing patches — recorded experiment

Read the figure

Connect it to the computation

1 Define patch layout and masked error

2 Separate visible features from full targets

3 Fit the decoder and audit the reconstruction

Run the example

Recorded reference execution

Make visible-output errors arbitrarily large

Make hidden information impossible to recover

Make it yours

Expose the leakage shortcut

Audit every patch of a rectangular color image

Check your understanding

Diagnose the result

Carry forward

Keep your evidence

Primary references