Cross-modal: connect images and text

Stage 1 / 8

Define the data contract

Begin with what one example means. Keep stable IDs and split before fitting normalization, tokenization or augmentation. A small synthetic fixture makes bugs visible; a separate documented dataset is needed to study generalization.

Keep pair and group IDs through every shuffle. Split related images and captions together; record the policy for multiple valid captions or near-duplicate images.

Practice in project stage 1

Available preparation

• From NumPy to jax.numpy
• Shapes, broadcasting, and dtypes
• Design an input pipeline

Stage 2 / 8

Build a transparent baseline

First build a model small enough to inspect. Predict intermediate shapes and check a forward pass using a separate calculation. Your baseline provides a reference for every later optimization.

Start with two small encoders and a declared embedding normalization rule. Trace both batch axes in the similarity matrix and compare one dot product independently.

Practice in project stage 2

Available preparation

• Build a tiny multilayer perceptron
• Model and state with Flax NNX

Stage 3 / 8

Make each update explainable

Record whether a reported loss belongs to the parameters before or after an update. Keep validation outside the training transition and save the configuration with the random seed. A falling loss shows fitting; it is verified separately from usefulness on new examples.

Use a stable contrastive objective with an explicit temperature and positive-pair mask. Test aligned versus deliberately permuted pairs; duplicate positives must not silently become negatives. The training-methods project supplies separate objective labs for these methods; follow each phase prerequisite before attempting them. Integrating every objective into this modality harness remains an extension.

Practice in project stage 3

Available preparation

• Gradient checking and numerical accuracy
• Optimize with Optax
• A compiled train and evaluation step
• Debug unstable learning
• Track experiments and lineage with MLflow
• Contrastive learning: views, positives and negatives
• LoRA: adapt, save and merge low-rank updates

Stage 4 / 8

Prove that training can resume

Weights alone cannot reconstruct an interrupted experiment. Save optimizer memory, random state, completed step and input iterator position together with the data and package contract. Compare the next several updates after restore.

Save both encoders, optimizer state and the paired sampler position together. Reject a checkpoint that restores only one side of the learned embedding space.

Practice in project stage 3

Available preparation

• Load, batch, and checkpoint order with Grain
• Save model, optimizer, and random state
• Recover data position and resume

Stage 5 / 8

Measure quality and inspect errors

Freeze evaluation inputs before tuning. Aggregate sums and counts instead of averaging unequal batch means. Show representative failures alongside a metric so a reader can see what the model gets wrong.

Evaluate image-to-text and text-to-image retrieval on a frozen candidate set. Report recall at declared ranks, candidate count and multiple-positive handling; changing the candidate pool changes the task.

Practice in project stage 3

Available preparation

• A compiled train and evaluation step
• Train a CNN on a small dataset
• Regularization and honest validation

Stage 6 / 8

Choose weights, activations and accumulation separately

Start from a float32 reference. Weight storage, activation computation, accumulation and optimizer memory are separate decisions. Integer quantization also needs scales, zero-point policy, calibration data and runtime support. A smaller file is not evidence of faster inference.

Compare encoder precision and stored embedding precision separately. Track rank changes and recall, not only elementwise embedding error.

Practice in project stage 4

Available preparation

• A Transformer block and mixed precision
• Choose weight, activation and accumulation precision

Stage 7 / 8

Verify the deployment boundary

Package preprocessing with the trained computation. Check exported outputs on fixed normal and boundary inputs before serving them. Framework conversion and copying weights between model definitions need explicit shape, layout and dtype contracts.

Version the two encoders, preprocessing assets and search index together. Refuse an index built by an incompatible encoder version. Use the weight-conversion project to audit an actual PyTorch/Flax dense–normalization–activation architecture. Extend its layerwise checks for this modality rather than assuming that its MLP mapping covers all operators.

Practice in project stage 4

Available preparation

• Move models between JAX, Keras, TensorFlow and PyTorch
• Export a computation and verify its serving contract
• Deploy at the edge: conversion, budgets and device checks
• Containerize a model service and verify its boundary
• Convert PyTorch weights to Flax and locate numerical errors

Stage 8 / 8

Measure inference and rehearse failure

Define the workload before timing it: input sizes, batching, concurrency and device. Warm up, synchronize device work and distinguish startup from steady execution. Test a failed load and a rollback to the last verified artifact.

Measure query encoding and search separately, then end-to-end response time. Test missing or corrupted inputs and rebuild/rollback of an index. Generative multimodal models are a later extension.

Practice in project stage 5

Available preparation

• Benchmark asynchronous work correctly
• Diagnose recompilation and host synchronization
• MLOps: data contracts, CI gates and monitoring
• ModelOps: own, approve, roll back and retire a model

Write a precision policy, not just a dtype

Record weights, activations, accumulation, optimizer state and any inference cache separately. Keep a float32 reference and compare the same held-out examples before choosing a lower-precision artifact.

• Float32: establish the reference output and quality first.
• Bfloat16 or float16: test numerical range and reduction behavior on the actual backend; retain sensitive reductions or optimizer state in higher precision where needed.
• Int8: document calibration examples, scale granularity, clipping and activation policy. Weight-only and weight-and-activation quantization are different experiments.
• Int4 or float8 extensions: require a supported runtime and hardware, declared packing/scaling rules and new quality measurements. No general compatibility or speedup is assumed.

For every supported policy keep artifact size, task quality, output error, latency and the actual device. Record unsupported policies as unsupported.

Cross-framework work begins with parity

Use the existing Keras/JAX and TensorFlow conversion labs as preparation. A model definition, a saved checkpoint and an inference artifact have different contracts. Moving weights does not automatically move optimizer state, preprocessing or training behavior.

• Pin both environments and choose a conversion path supported by those versions.
• Map parameter names, kernel layouts, input shapes and dtypes explicitly.
• Compare intermediate tensors and final outputs on fixed inputs, including boundary cases.
• Compare task quality after export and again after quantization; measure the exported runtime itself.

PyTorch, Keras, TensorFlow and JAX interoperability is tested per model and operation set. A successful tiny example is verified separately from that every model can round-trip between them.

Grow the harness after the small version works

Share run configuration, metrics, checkpoint metadata and artifact packaging. Keep modality-specific preprocessing, objectives and decoding behind explicit interfaces. Extend one working image system first, then text, audio and paired retrieval.

MaxText is an advanced reference for language-model training and supported multimodal workflows. Study its configuration, input pipeline, checkpointing, sharding and precision choices. Pin a release or commit in a separate compatible environment before adapting code; the course CPU environment does not validate a MaxText installation.

Carry engineering evidence across the lifecycle

Use the engineering-release project after the modality harness. Track the same data, preprocessing, model and evaluation identity from experiment through service.

• MLflow: compare runs with the same metric definition and preserve data/code lineage.
• Containers: record base/image/model digests, non-root runtime, mounted artifacts and readiness.
• MLOps: validate data contracts, automate repeatable evaluation, inspect slices and monitor delayed quality.
• LLMOps: version prompts, retrieval and tools; retain calibrated evaluation cases and redacted traces.
• ModelOps: assign ownership, bind approval to a specific bundle and rehearse rollback and retirement.

MLflow tracking and Docker are concrete tools; MLOps, LLMOps and ModelOps describe responsibilities around them. Read the linked operations lessons before treating a registry alias or a passing local check as a deployed system.

References and next steps

• MaxText: advanced language-model training reference
• MaxText multimodal tutorial
• MaxDiffusion: image-generation extension reference

Open the connected harness

Open the preparation project

Open the preparation project

Open the preparation project

Planned extensions

These project specifications build on this track. Their workspaces and checks are not available yet.

• Train a CLIP image–text model — Learn paired representations, test zero-shot classification, and qualify both retrieval directions before releasing the encoders.
• Train and compare a SigLIP model — Replace the softmax objective with pairwise sigmoid training. Audit labels, reduction, negatives and distributed gradients against a controlled CLIP baseline.
• Build a measured retrieval service — Take a qualified text or cross-modal encoder through exact search, approximate indexes, hybrid retrieval, reranking, evaluation and safe index migration.
• Build a grounded multimodal model — Qualify encoders, connector alignment, instruction adaptation and grounded evaluation, keeping retrieval a distinct outcome.