Blend out-of-fold predictions only over members that predicted each row

[fus3r]

Fixes #55.

The bug

When max_num_rows is set, each ensemble member subsamples a different set of rows. So in predict_oof_proba a member only produces out-of-fold predictions for its own rows, and the rest of its output stays zero.

fit() builds the calibration input by averaging over all members, zeros included. That makes rows sum to k / n_estimators (k = how many members actually predicted the row) instead of 1, so the calibrator is fit on vectors that aren't valid probabilities. You only hit it with max_num_rows, which is the large-data case.

The single-split path made it worse: it kept only member 0's validation indices and sliced every member with them. The multi-fold path averaged over all N rows, including ones a member never sampled.

enable_nnls can't be combined with max_num_rows (the constructor rejects it), so the NNLS blend never runs into this and I left it as is.

Repro on main

Random weights, PyTorch, no checkpoint needed:

import numpy as np, torch
from tabfm.src.pytorch import model as pytorch_model
from tabfm.src.classifier_and_regressor import TabFMClassifier

torch.manual_seed(0)
model = pytorch_model.TabFM(
    embed_dim=8, max_classes=3, col_num_blocks=1, col_nhead=2, col_num_inds=8,
    row_num_blocks=1, row_nhead=2, row_num_cls=2, icl_num_blocks=1, icl_nhead=2,
    ff_factor=2, feature_group_size=2, is_classifier=True)

np.random.seed(0)
X, y = np.random.rand(40, 3), np.random.randint(0, 2, 40)

captured = {}
orig = TabFMClassifier._fit_calibration
def spy(self, P, y_):
    captured["P"] = np.asarray(P).copy()
    return orig(self, P, y_)
TabFMClassifier._fit_calibration = spy

clf = TabFMClassifier(
    model=model, n_estimators=3, batch_size=2, random_state=42, max_num_rows=30,
    binary_calibration_method="platt", num_folds_for_cv=2,
    min_rows_for_single_val_split=5)
clf.fit(X, y)
print(np.round(captured["P"].sum(axis=1), 3))

On main many of those sums come out as 0.333 or 0.667. With the fix they are all 1.0.

Fix

predict_oof_proba now stores which rows each member predicted (oof_valid_mask_). fit() drops rows that no member predicted, and for the rest it averages each row over just the members that predicted it. Without max_num_rows every member predicts every row, the mask is all True, and the result is unchanged.

Tests

Two tests added: one checks the mask matches each member's predicted rows under subsampling, the other checks the calibration input equals the per-row mean over the predicting members (with a never-predicted row dropped and the labels realigned). pytest tabfm/src/ is green (74 tests).

[fus3r]

Hi, just following up on this PR. Is there anything you'd like me to add or change before review?

[erzel]

Thank you for the contribution and apologies for the delay in reviewing.
We have an internal fix for this which we plan to publish shortly which uses a shared hold-out set instead of CV predictions, instead of fitting calibration on partial ensembles. We believe that would be a better solution since it allows calibration fitting to be closer to inference (where all ensemble members are used not just a subset).

[fus3r]

Thanks for the update, and no worries about the delay. The hold-out approach sounds like the cleaner way to handle it. Happy to try the fix when it lands!

[erzel]

Sorry for the long delay, but the hold-out set fix is now merged (PR #85). Let us know if you have any more issues.

Thank you for your patience.