Skip to content

[BUG] pretrain() -> fit() discards everything pretraining learned #446

Description

@ChrisW09

Describe the bug

The documented warm-start workflow does not work. docs/api/training/index.rst states:

model.build_model(...); model.pretrain(...); model.fit(...) ... "pretrain() updates the live
model's embeddings, so the following fit() continues from the pretrained weights."

But fit() defaults to rebuild=True, which calls _build_model() and constructs a brand-new
TaskModel/architecture — throwing the pretrained weights away.

Verified independently with lr=1e-12 (so training cannot move any weight): embedding tensors after
fit() differ from the pretrained ones by up to 4.29 in absolute value. With rebuild=False the
max difference is exactly 0.0.


pretrain() → fit() throws the pretrained embeddings away: fit() defaults to rebuild=True and re-initialises the backbone

Where: deeptab/models/_mixins/fit.py (301, 445-469)

The documented warm-start workflow (build_model(); pretrain(); fit()) does not work — fit() defaults to rebuild=True, which calls _build_model() and constructs a brand-new TaskModel/architecture, discarding everything pretrain() learned.

Observed: With lr=1e-12 (training cannot move any weight) the embedding tensors after fit() differ from the pretrained ones by up to 4.4 in absolute value, e.g. {'embedding_layer.num_embeddings.0.0.weight': 0.464, 'embedding_layer.cat_embeddings.0.0.weight': 4.417}. id(m._task_model.estimator) also changes across the fit(). Passing rebuild=False gives max diff 0.0 for every tensor. (With mismatched build/fit split seeds the rebuilt embeddings can even change shape: 16 vs 14.)

Expected: docs/api/training/index.rst:90-105 states: "model.build_model(...); model.pretrain(...); model.fit(...)" and "pretrain() updates the live model's embeddings, so the following fit() continues from the pretrained weights." The pretrained embedding weights should survive the subsequent fit().

Repro
import warnings; warnings.simplefilter('ignore')
import copy, numpy as np, pandas as pd
from deeptab.models.fttransformer import FTTransformerClassifier
from deeptab.configs import FTTransformerConfig, TrainerConfig
rng = np.random.default_rng(0)
X = pd.DataFrame({'a': rng.normal(size=64), 'b': rng.normal(size=64), 'g': rng.choice(list('xyz'), 64)})
y = (X['a'] > 0).astype(int).values
TC = TrainerConfig(max_epochs=1, batch_size=32, val_size=0.2, patience=3, lr=1e-12)  # lr~0: weights cannot move
m = FTTransformerClassifier(model_config=FTTransformerConfig(d_model=16, n_layers=1, n_heads=2),
                            trainer_config=TC, random_state=0)
m.build_model(X, y, batch_size=32, val_size=0.2, random_state=0)
m.pretrain(pretrain_epochs=2, k_neighbors=3, save_path='pre.pth')
pre = copy.deepcopy(m._task_model.estimator.get_embedding_state_dict())
m.fit(X, y, accelerator='cpu', devices=1)          # default rebuild=True
post = m._task_model.estimator.get_embedding_state_dict()
print({k: float((pre[k]-post[k]).abs().max()) for k in pre})

Expected behavior
Either fit() after pretrain() should default to continuing from the pretrained weights, or
pretrain() should record that a rebuild would discard its work and warn/raise. The docs promise the
former.

Screenshots
n/a

Desktop (please complete the following information):

  • OS: macOS (Darwin 25.5.0, arm64)
  • Python version: 3.11.15
  • deeptab Version: 2.0.0 (main @ 4e6a359)

Additional context
torch 2.9.1, lightning 2.6.5, scikit-learn 1.9.0, numpy 2.4.6. Found in a second-pass review of v2.0.0
(seven independent lenses, each finding adversarially re-verified by a second reviewer, then re-run by
hand). Distinct from the already-filed #409-#426.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions