[BUG] encode() returns dropout-corrupted embeddings and is completely broken for LSS models #451

Describe the bug

Defects in the public encode() API.

encode() never switches the model to eval mode, so embeddings after fit() are dropout-corrupted and non-deterministic

Where: deeptab/models/_mixins/predict.py (126-170)

predict()/predict_proba() call self._task_model.eval() first, but encode() does not; fit() leaves the module in train mode, so encode() runs the backbone with dropout active and returns a different tensor on every call.

Observed: training mode: True; two consecutive encode() calls on the same rows differ by up to 1.47 (identical: False). After a manual _task_model.eval() the calls agree exactly, and differ from the train-mode output by up to 1.34. Consequently fitted.encode(X) and Model.load(path).encode(X) (load() does call eval()) never match.

Expected: encode() is documented as "Return dense embedding vectors from the model backbone" — a deterministic feature-extraction call. It should put the model in eval mode like predict() does.

import warnings; warnings.simplefilter('ignore')
import torch, numpy as np, pandas as pd
from deeptab.models.fttransformer import FTTransformerRegressor
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'].values * 2
m = FTTransformerRegressor(model_config=FTTransformerConfig(d_model=16, n_layers=1, n_heads=2),
                           trainer_config=TrainerConfig(max_epochs=1, batch_size=16, patience=3), random_state=0)
m.fit(X, y, accelerator='cpu', devices=1)
print('training mode:', m._task_model.estimator.training)
a, b = m.encode(X.iloc[:10]), m.encode(X.iloc[:10])
print('identical:', torch.allclose(a, b), float((a-b).abs().max()))
m._task_model.eval()
c = m.encode(X.iloc[:10]); print('after eval identical:', torch.allclose(c, m.encode(X.iloc[:10])))

SklearnBaseLSS.encode() unpacks 2 values from 3-element batches, so it always raises on every LSS model

Where: deeptab/models/lss_base.py (510-512)

for num_features, cat_features in tqdm(data_loader) assumes the prediction dataset yields (num, cat) pairs, but preprocess_new_data yields (num_features, cat_features, embeddings) triples. The loop raises ValueError: too many values to unpack (expected 2) on the first batch, so encode() can never succeed on any *LSS estimator. (Even if it unpacked correctly, the next line estimator.encode(num_features, cat_features) would bind cat_features to BaseModel.encode's grad parameter, since that signature is encode(self, data, grad=False).)

Observed: ValueError: too many values to unpack (expected 2)

(The same architecture wrapped in FTTransformerClassifier — which uses _PredictMixin.encode instead of this override — encodes fine: shape (32, 4, 16).)

Expected: SklearnBaseLSS.encode should iterate the 3-tuple batch and pass it as a single data argument, like _PredictMixin.encode does (emb = self._task_model.estimator.encode(batch)). The method has a full numpydoc contract documenting a returned tensor, and no LSS model can honour it.

import warnings; warnings.simplefilter("ignore")
import numpy as np, pandas as pd
from deeptab.models import FTTransformerLSS
from deeptab.configs import FTTransformerConfig, TrainerConfig

rng = np.random.default_rng(0)
X = pd.DataFrame(rng.normal(size=(32, 4)), columns=list("abcd"))
y = X.values @ rng.normal(size=4)

m = FTTransformerLSS(model_config=FTTransformerConfig(d_model=16, n_layers=1, n_heads=2),
                     trainer_config=TrainerConfig(max_epochs=1, batch_size=16,
                                                  patience=1, lr_patience=1))
m.fit(X, y, family="normal", accelerator="cpu",
      enable_progress_bar=False, enable_model_summary=False)
m.encode(X)

The encode() docstring example uses MLPClassifier and promises a 2-D result; MLP can never encode and successful models return 3-D

Where: deeptab/models/_mixins/predict.py (140-158 (Returns block at 142-143, Examples at 152-157))

_PredictMixin.encode documents Returns: torch.Tensor of shape (n_samples, embedding_dim) and an example >>> clf = MLPClassifier(); clf.fit(...); clf.encode(X_test) → torch.Size([100, 64]). Neither holds. BaseModel.encode requires an embedding_layer and one of mamba/rnn/lstm/encoder; MLP has neither, so the documented example always raises ValueError: The model does not have an embedding layer. For models that do work, the return is the un-pooled token sequence (n_samples, n_features, d_model) — 3-D, not 2-D. docs/getting_started/faq.md:629 repeats the same wrong shape ("returns dense representations as a tensor of shape (n_samples, embedding_dim)").

Observed: MLPClassifier.encode -> ValueError The model does not have an embedding layer
FTTransformer.encode shape: (32, 4, 16) # 3-D (n_samples, n_features, d_model), not (n_samples, embedding_dim)

Expected: The example should use an architecture that actually supports encoding (FTTransformer/SAINT/TabTransformer/TabulaRNN), the Returns section should state the true 3-D shape (or the method should pool to 2-D), and both this docstring and faq.md:627-637 should say which model families support encode() — the FAQ currently presents it as universally available on "a fitted model".

import warnings; warnings.simplefilter("ignore")
import numpy as np, pandas as pd
from deeptab.models import MLPClassifier, FTTransformerClassifier
from deeptab.configs import FTTransformerConfig, TrainerConfig

rng = np.random.default_rng(0)
X = pd.DataFrame(rng.normal(size=(32, 4)), columns=list("abcd"))
y = rng.integers(0, 2, 32)
tc = TrainerConfig(max_epochs=1, batch_size=16, patience=1, lr_patience=1)
kw = dict(accelerator="cpu", enable_progress_bar=False, enable_model_summary=False)

clf = MLPClassifier(trainer_config=tc); clf.fit(X, y, **kw)
try:
    clf.encode(X)                       # the docstring's own example
except Exception as e:
    print("MLPClassifier.encode ->", type(e).__name__, e)

ft = FTTransformerClassifier(model_config=FTTransformerConfig(d_model=16, n_layers=1, n_heads=2),
                             trainer_config=tc)
ft.fit(X, y, **kw)
print("FTTransformer.encode shape:", tuple(ft.encode(X).shape))

Expected behavior
encode() should call .eval() (and run under torch.no_grad()) so embeddings are deterministic and
dropout-free, and the LSS override should unpack batches the same way the base implementation does.

Related: the stray self.encoder(x) line in the same method is already filed as #426 item 1.

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.