Pretraining And Fine-Tuning

Pretraining learns from relationships in a broad collection of records. Fine-tuning starts from those learned weights and adapts the model to a specific task. In relflow, both use trainer.fit(...); the data and reconstruction objectives determine what the model learns.

Stage	Data	Objectives
Self-supervised pretraining	Records without task labels	Reconstruct sampled observed values with `rf.Mask(rate=..., reconstruct=True)`.
Supervised fine-tuning	Labeled task records	Predict always-hidden labels with `mask=True`, optionally retaining some reconstruction objectives.
Resume a run	The same training task	Restore the model, optimizer, scheduler, and loop progress from a full Lightning checkpoint.

Pretraining is optional. A supervised model can also start from newly initialized weights. An ordinary float mask, such as mask=0.15, only supplies input dropout; pretraining needs reconstruct=True to create an objective.

Pretrain On Histories

This customer record has no churn label. Purchase amounts and merchant types provide their own supervision when selected values are hidden:

tenure_days: 420
region: northeast
purchases:
  - days_ago: 12
    amount: 42.75
    merchant: grocery
  - days_ago: 3
    amount: 18.50
    merchant: cafe

import lightning.pytorch as lit

import relflow as rf

model = rf.Model.md(
    tenure_days=rf.Number,
    region=rf.Category,
    purchases=rf.Branch(
        length=64,
        overflow="tail",
        days_ago=rf.Number,
        amount=rf.Number(mask=rf.Mask(rate=0.15, reconstruct=True)),
        merchant=rf.Category(
            mask=rf.Mask(rate=0.15, reconstruct=True),
        ),
    ),
)
model.optimizer = rf.adamw(learning_rate=1e-3)
pretraining_data = rf.ArrowDataModule(
    model=model,
    train=pretrain_table,
    validate=pretrain_validation_table,
)
pretrainer = lit.Trainer(max_epochs=50)
pretrainer.fit(model=model, datamodule=pretraining_data)
model.save("artifacts/customer-pretrained.rf")

The two tables are independent, application-supplied Arrow splits with this record structure. Histories are sorted oldest to newest for overflow="tail". Use rf.PolarsDataModule with eager frames when the application prepares its splits in Polars. Keep preparation consistent through the preprocessing pipeline.

Each sampled value is predicted from the remaining context; unselected values remain inputs. These fractional reconstruction policies also sample during validation and test, but are inactive during ordinary prediction. See Masking for selection and stage rules.

The example saves the final model. Use the checkpoint selection described in Training to save the best validation state instead.

Evaluate Vocabulary Coverage

For large-vocabulary reconstruction, report how many selected answers the model can express as well as how well it predicts them. For example, monitor .purchases.merchant/validate.coverage.content alongside .purchases.merchant/validate.nll.content and .purchases.merchant/validate.accuracy.all. The Category metrics and Cluster metrics define these denominators and count unavailable targets explicitly.

An unknown Category or Cluster target contributes no label gradient. Forcing its prediction to be uniform over known labels would teach a different task, not recover the missing answer. With no known targets, conditional reconstruction quality is NaN; zero coverage means the validation objective is unmeasurable, not that the model reconstructed everything perfectly. Cluster balancing and adaptive-cluster updates also exclude unknown answers. Do not select a checkpoint from such a conditional score without checking coverage.

Vocabulary growth can change the evaluated known-target population between epochs. NLL can increase when newly admitted, difficult labels join it even if predictions on the old labels improve. The summed validation loss is not a coverage-adjusted model-quality score. Keep a fixed validation source, monitor coverage, and compare known-target quality at comparable coverage. If necessary, retain a fixed cohort of training-established labels for an additional stable diagnostic; do not populate the model vocabulary from validation or test data.

Random reconstruction masks introduce another source of variation. For a repeatable diagnostic, supply application-owned Boolean selection fields and use Mask(query=..., reconstruct=True). This fixes selected coordinates; it does not freeze vocabulary growth or make unknown answers learnable. Keep any held-out calibration set separate from model selection, and evaluate downstream confidence separately for familiar and unseen input identities.

P062 isolates vocabulary coverage with fixed query-selected masks and checks that rebatching the same targets cannot change the reported counts or conditional scores. P065 tests Cluster coverage.

P066 separately checks cold Category inputs across successive training cohorts. It currently exposes a remaining failure: admitting an input identity can change reconstruction NLL without an optimizer step, even at complete target coverage. Target coverage does not measure input familiarity. Keep input cohorts separate when assessing pretraining stability; fixed masks alone cannot remove this source of drift. Initialization alternatives that pass the admission check have failed other learning proofs and are not enabled in the model.

Set metrics count positive memberships rather than scalar targets. A partial-OOV Set retains its known positive and negative supervision; dropping the whole set would discard valid evidence. An all-OOV set still supplies known negatives when annotations are complete. Report member coverage, complete-set coverage, exact-set accuracy, and micro F1 together. Empty sets are valued, complete annotations, although an all-empty evaluation has no positive-member denominator. These semantics are not appropriate for incompletely annotated labels without further masking.

Adapt To A Labeled Task

For fine-tuning, each customer record gains churned: false or churned: true, defined over a future outcome window. Its purchase history ends at the decision time. Pretraining and fine-tuning splits must respect the same evaluation boundaries; pretraining on held-out customers or future observations changes what the downstream evaluation measures.

Load the pretrained artifact, remove its sampled reconstruction policies, and add a supervised label:

finetuned = rf.Model.load("artifacts/customer-pretrained.rf")
reconstructed_inputs = rf.where("address").is_in(
    [rf.Address("purchases", "amount"), rf.Address("purchases", "merchant")]
)
finetuned.update(reconstructed_inputs, mask=False)
finetuned.extend(
    rf.where("address") == rf.Address(),
    churned=rf.Boolean(mask=True),
)

Amount and merchant remain visible inputs. Churn is always excluded from input embedding, while its observed label supplies the training loss. If the target already exists in the schema, use finetuned.update(target, mask=True) instead of adding it again.

Customer context combines tenure, region, and repeated purchases containing age, amount, and merchant. Fine-tuning adds a churn prediction at the customer root. — Figure 1: Pretraining reconstructs selected amounts and merchants. Fine-tuning reuses the customer hierarchy and adds a churn target.

Schema edits rebuild modules and retain state entries whose names, types, and tensor shapes still match. The new target starts with newly initialized state; changing widths, field types, or structure can also reinitialize incompatible parts. Keep existing addresses and geometry stable when reusing a pretrained model. Adding a field can change context and reduction behavior even when existing weights transfer. Schema Mutation describes these boundaries.

Start A New Training Run

Assign a fresh optimizer factory after the edits and use a new Trainer with labeled Arrow splits. A smaller learning rate is a starting point to evaluate:

finetuned.optimizer = rf.adamw(learning_rate=1e-4)
finetuning_data = rf.ArrowDataModule(
    model=finetuned,
    train=finetune_table,
    validate=finetune_validation_table,
)
finetuner = lit.Trainer(max_epochs=20)
finetuner.fit(model=finetuned, datamodule=finetuning_data)
finetuned.save("artifacts/customer-churn.rf")

This starts new optimization from the transferred model state. Use ckpt_path with a full Lightning checkpoint when resuming the same task, including its optimizer and loop progress. Reapply any processors, scheduler, and desired compilation after loading and schema edits. For DDP, every rank must construct the same final schema.

Fine-tuning here updates the model end to end, including the new decoder and existing encoders. Loaded vocabularies and Number normalization statistics also continue learning from training observations; loading does not freeze them. Validation, test, and prediction reuse that state. For a fixed feature extractor, set embed=True on the desired root or branch, export embeddings in prediction mode, and fit a separate downstream model on those vectors.

At inference, omit churned from the input and read the prediction at /churned. See Prediction Output.

Keep Auxiliary Objectives Deliberately

To retain amount reconstruction alongside churn, keep that policy when adapting the pretrained artifact. Replace the earlier finetuned.update(reconstructed_inputs, mask=False) call with:

finetuned.update(
    rf.where("address") == rf.Address("purchases", "merchant"),
    mask=False,
)
finetuned.update(
    rf.where("address") == rf.Address("purchases", "amount"),
    weight=0.1,
)

Then add the churn target and start the new run as above. Keeping the amount policy preserves its trained decoder; removing reconstruction and later reenabling it creates a new decoder. Both amount and churn losses are now optimized. weight scales the amount loss rather than its sampling rate or a fixed fraction of the combined loss.

Compare fine-tuning against the same supervised task trained from scratch. Use the same held-out data and churn metric: pretraining loss and fine-tuning loss have different objective sets and are not directly comparable. Retaining auxiliary objectives may help or hurt the task, so evaluate that choice rather than assuming pretraining guarantees a better model. See Evaluation.