google-research / google-research/tabfm
`max_num_rows` + calibration: single-split path fits calibration on partially-zeroed OOF probabilities
- 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
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