Skip to content

Foundation Model API

The foundation model learns general-purpose entity representations from behavioral event data using self-supervised learning. It is configured via a YAML file and trained using the pretrain() function.

pretrain()

Combines behavioral representation fitting and foundation model training into a single call.

from monad.ui import pretrain

pretrain(
    config_path=Path("fm_config.yaml"),
    output_path=Path("./foundation_model"),
)

Signature

def pretrain(
    config_path: Path,
    output_path: Path,
    use_last_basket_sketches: bool = True,
    recency_sketch_timespan_days: Optional[int] = 60,
    storage_config_path: Optional[Path] = None,
    uniqueness_threshold: float = 0.9,
    nan_threshold: float = 0.9,
    sketch_depth: Optional[int] = None,
    sketch_width: Optional[int] = None,
    callbacks: Optional[list[pytorch_lightning.Callback]] = None,
    pl_logger: Optional[pytorch_lightning.loggers.Logger] = None,
    resume: bool = False,
    overwrite: bool = False,
    seed: Optional[int] = None,
    quick_check: bool = False,
    ignorable_modalities: Optional[dict[str, set[str]]] = None,
) -> None

Parameters

Parameter Type Default Description
config_path Path required Path to the YAML pretraining config file.
output_path Path required Directory where all output artifacts (checkpoints, features, logs) will be stored.
use_last_basket_sketches bool True Whether to add a sketch with event data from the immediate past as a separate input to the model.
recency_sketch_timespan_days int \| None 60 Timespan in days for recency sketches. Set to None to disable recency sketches.
storage_config_path Path \| None None Path to the storage configuration file (credentials for cloud filesystems).
uniqueness_threshold float 0.9 Maximum fraction of unique values for a categorical column. Columns above this threshold are excluded.
nan_threshold float 0.9 Maximum fraction of NaN values for a column. Columns above this threshold are excluded.
sketch_depth int \| None None Sketch depth override (for testing). If not set, determined automatically.
sketch_width int \| None None Sketch width override (for testing). If not set, determined automatically.
callbacks list[Callback] \| None None PyTorch Lightning callbacks to attach to the trainer.
pl_logger Logger \| None None PyTorch Lightning logger instance (e.g., MLFlowLogger, TensorBoardLogger).
resume bool False Whether to reuse existing partial results from a previous interrupted run.
overwrite bool False Whether to overwrite all existing data in the output directory and start fresh.
seed int \| None None Random seed for reproducibility. If not provided, reproducibility is not guaranteed.
quick_check bool False Validate configuration, connection, and schema wiring on sampled data with a shortened training run. Outputs are marked _QUICK_CHECK and are not suitable for production or downstream use. Use normal capped training for a downstream-ready smoke checkpoint.
ignorable_modalities dict[str, set[str]] \| None None Per-data-source set of column names to flag as ignorable. See Ignorable Modalities below.

Ignorable Modalities

Marking a modality as ignorable makes it eligible for probabilistic sketch dropout during foundation model training: on each training sample its sketch is randomly replaced with the default (empty) sketch. This acts as regularization — the model learns not to over-rely on that modality and stays robust when the column is sparse or missing at inference time.

Pass a mapping from data source name to the set of column names to flag:

from pathlib import Path
from monad.ui import pretrain

pretrain(
    config_path=Path("fm_config.yaml"),
    output_path=Path("./foundation_model"),
    ignorable_modalities={
        "transactions": {"discount_code", "campaign_id"},
        "web_events":   {"referrer"},
    },
)

Dropout probability

The per-sample dropout is controlled by the IGNORABLE_SKETCH_DROPOUT_PROBABILITY environment variable (default 0.5); a modality's sketch is dropped when random() > IGNORABLE_SKETCH_DROPOUT_PROBABILITY. Columns not listed are never dropped. This is the pretrain-stage counterpart of ignored_modalities at the scenario stage, which drops a modality always rather than probabilistically.

Two-Stage Functions

The pretrain() function internally runs two stages. You can also call them individually for more control:

fit_behavioral_representation()

First stage: fits the behavioral representation (feature engineering, sketch computation).

from monad.run import fit_behavioral_representation

fit_behavioral_representation(
    config_path=Path("fm_config.yaml"),
    output_path=Path("./foundation_model"),
)

train_foundation_model()

Second stage: trains the neural foundation model using results from the fit stage.

from monad.run import train_foundation_model

train_foundation_model(
    output_path=Path("./foundation_model"),
    callbacks=[...],
    pl_logger=mlflow_logger,
)

Note

train_foundation_model() must be run after fit_behavioral_representation(). Both must use the same output_path.

Output Structure

After successful training, the output directory contains:

foundation_model/
├── fm/                      # Foundation model checkpoint
│   ├── best_model.ckpt
│   ├── lightning_checkpoints/
│   └── _FINISHED            # Empty marker file: the checkpoint is complete
├── config.yaml              # Resolved config
├── suggested_config.yaml    # Config with column report suggestions applied
└── features/                # Pre-computed behavioral features

See Also