`balance_dataset` can make the dataset less balanced

[itrujnara]

I am working on a difficult classification problem, with very noisy and strongly unbalanced data. Since data collection is slow (it is clinical data), I decided to try data simulation as an alternative approach to improve the models. TabPFGen is my best shot, as I have multiple correlations between features that need to be accounted for.

I have tried both generate_classification and balance_dataset, and neither achieved the desired effect. In fact, the balancing made the problem worse, as after refining nearly all generated samples were assigned to the majority class.

Minimum working example (sort of)

I cannot share my data for privacy reasons, and I have not been able to reproduce the issue fully even with the most heinous fake data. The script below is as close as I could get, with about 30% of the generated samples assigned to the majority class.

from sklearn.datasets import make_classification
from tabpfgen import TabPFGen

X, y = make_classification(
    n_samples=3000, n_classes=2,
    n_features=400, n_informative=2,
    weights=[0.99, 0.01], random_state=42
)

generator = TabPFGen()

X_synth, y_synth, X_combined, y_combined = generator.balance_dataset(
    X, y
)

Observed behavior

With my training data, approximately 99% of the generated samples are assigned to the majority class, as demonstrated by the log:

=== Final Statistics ===
Final combined class distribution:
  Class 0: 4018 total (2956 original + 1062 synthetic)
  Class 1: 1894 total (44 original + 1850 synthetic)
Dataset balancing completed!

Expected behavior

I hoped that at least 20% of the synthetic samples would be assigned to the positive class.

Final remarks

I am not surprised by the behavior, as the data is particularly complex and all classifiers I have tried so far suffer from similar issues when applied with the default hyperparameters. It would, however, be nice to get slightly more balanced outputs, as this is one of the stated purposes of this package. All that said, thank you for all the work on this, as the model indeed behaves well with more generic data.