google-research / google-research/tabfm

`max_num_rows` + calibration: single-split path fits calibration on partially-zeroed OOF probabilities

Open
#55 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.6k
Forks
270
Avg merge
1d 7h
Merged PRs (30d)
1

Description

## Summary

When `max_num_rows` triggers row subsampling, a calibration method is enabled, and the single validation-split path is taken, the out-of-fold probability matrix passed to `_fit_calibration` contains rows that do not sum to 1. For example, with `n_estimators=3`, some rows sum to `0.333` or `0.667`.

The calibrator is therefore fit on a mixture of real predicted probabilities and zero-filled missing predictions, so calibration is silently wrong for exactly the large-dataset case where users are likely to enable `max_num_rows`.

Verified against `main` @ `65aeeed`.

## What's happening

Each ensemble member draws its own row subsample, so each member's validation rows correspond to different absolute row indices.

`predict_oof_proba` writes each member's OOF predictions to the correct places:

```python
outputs_oof[i, val_indices_list[i]] = out_i
```

But on the single-split path it only records member 0's validation indices:

```python
self.oof_val_indices_ = val_indices_list[0]
```

Then `fit()` slices all members using member 0's validation indices:

```python
oof_probs_fit = oof_probs[:, val_idx, :]
```

For members other than 0, most of those rows were never predicted and are still all-zero. Those zeros are included in:

```python
P = np.mean(oof_probs_fit, axis=0)
```

and then passed into `_fit_calibration`, with labels taken from member 0's validation rows.

## Repro

CPU, PyTorch backend, random-init weights. No checkpoint is needed because the issue is index bookkeeping in `fit()`. The spy only captures what `_fit_calibration` receives.

```python
import numpy as np
import 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 = np.random.rand(40, 3)
y = 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,
)

try:
clf.fit(X, y)
finally:
TabFMClassifier._fit_calibration = orig

print(np.round(captured["P"].sum(axis=1), 3))
```

Output:

```text
[0.667 1. 0.667 0.667 0.667 0.667 0.667 0.667 0.333 0.333 0.333 1.
0.667 0.333 1. ]
```

12 of 15 rows sum to `k / 3`, where `k` is the number of ensemble members that actually predicted that row.

## Expected behavior

Rows passed to `_fit_calibration` should be valid probability distributions, so each row of `P` should sum to approximately 1.

## Actual behavior

Rows passed to `_fit_calibration` can sum to less than 1 because zero-filled missing OOF predictions are averaged together with real predictions.

## Possible fix direction

The single-split path should either:

1. use a shared validation index set across ensemble members when calibration OOF predictions are needed, or
2. track valid OOF rows per estimator and average only over estimators that actually predicted each row, while aligning `y` to those absolute row indices.

At minimum, calibration should avoid averaging zero-filled missing predictions into `P`.

Contributor guide

Open the contributing guide

Research direction

Start in tabfm/src/classifier_and_regressor.py by tracing fit(), predict_oof_proba, and _fit_calibration; inspect how oof_val_indices_ and val_indices_list are used on the single-split path. Reproduce the issue with the supplied script, then verify that calibration receives aligned predictions whose rows sum to approximately 1 without including zero-filled missing outputs.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.