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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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