feat: add optional resampler= kwarg to AutoML.fit for per-fold class-imbalance handling (#1200)#1568
Open
immu4989 wants to merge 3 commits into
Open
feat: add optional resampler= kwarg to AutoML.fit for per-fold class-imbalance handling (#1200)#1568immu4989 wants to merge 3 commits into
immu4989 wants to merge 3 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds an opt-in resampler= hook to AutoML.fit to support per-fold resampling (e.g., SMOTE-style) inside FLAML’s CV/holdout evaluation loop, avoiding synthesized-sample leakage into validation folds as requested in issue #1200.
Changes:
- Add
resampler=Noneparameter toAutoML.fit, validate inputs, and store the resampler on the task for downstream access. - Apply per-fold resampling in
get_val_lossby cloning the provided resampler and callingfit_resample(X_train, y_train)beforeestimator.fit(...). - Add a new test module to exercise the new API and its validation behavior.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
flaml/automl/automl.py |
Introduces resampler argument and performs early validation / task attachment. |
flaml/automl/ml.py |
Applies the resampler per fold/call (clone + fit_resample) just before estimator training. |
test/automl/test_resampler.py |
Adds tests for the new resampler= behavior and validation. |
Comment on lines
1845
to
1849
| mlflow_logging=None, | ||
| fit_kwargs_by_estimator=None, | ||
| mlflow_exp_name=None, | ||
| resampler=None, | ||
| **fit_kwargs, |
Comment on lines
+2348
to
+2360
| if resampler is not None: | ||
| if "sample_weight" in fit_kwargs: | ||
| raise ValueError( | ||
| "Cannot combine 'resampler' with 'sample_weight' — resampling breaks " | ||
| "the 1-to-1 row alignment with sample weights. Use either resampling " | ||
| "or sample weighting, not both." | ||
| ) | ||
| if not hasattr(resampler, "fit_resample"): | ||
| raise TypeError( | ||
| "'resampler' must expose a fit_resample(X, y) -> (X, y) method " | ||
| "(e.g., an imbalanced-learn BaseSampler such as SMOTE)." | ||
| ) | ||
| task._resampler = resampler |
Comment on lines
+61
to
+63
| imblearn = pytest.importorskip("imblearn.over_sampling") | ||
| SMOTE = imblearn.SMOTE | ||
|
|
Comment on lines
+84
to
+110
| def test_resampler_leaves_validation_untouched(): | ||
| """Sanity check: the CV validation partitions must retain the raw class | ||
| distribution. If SMOTE were leaking into the validation folds, the search | ||
| would perceive an artificially balanced eval set and the val_loss reported | ||
| by the resampled fit would be systematically better than what the same | ||
| model achieves on the raw distribution. | ||
|
|
||
| We approximate this by asserting the final CV val_loss for the resampled | ||
| fit is not negative (it is 1 - f1, which is bounded in [0, 1] on a raw | ||
| imbalanced validation set with a non-trivial model). | ||
| """ | ||
| imblearn = pytest.importorskip("imblearn.over_sampling") | ||
| SMOTE = imblearn.SMOTE | ||
|
|
||
| X, y = _imbalanced_dataset(seed=1) | ||
| resampled = AutoML() | ||
| resampled.fit( | ||
| X_train=X, | ||
| y_train=y, | ||
| resampler=SMOTE(random_state=1, k_neighbors=3), | ||
| seed=1, | ||
| **_fit_settings(), | ||
| ) | ||
| assert 0.0 <= resampled.best_loss <= 1.0, ( | ||
| f"best_loss ({resampled.best_loss}) outside expected [0, 1] range for 1-f1 on a " | ||
| "raw imbalanced validation fold — validation may have been resampled" | ||
| ) |
Comment on lines
+52
to
+81
| def test_resampler_changes_chosen_config(): | ||
| """Passing a resampler should influence the search — the chosen best_config | ||
| on an imbalanced dataset with SMOTE will differ from the same fit without. | ||
|
|
||
| This is a proxy for verifying the per-fold hook actually fires; if the | ||
| hook were a no-op, both fits would land on the same best_config since | ||
| everything else about the search is deterministic (same seed, same | ||
| estimator, same data). | ||
| """ | ||
| imblearn = pytest.importorskip("imblearn.over_sampling") | ||
| SMOTE = imblearn.SMOTE | ||
|
|
||
| X, y = _imbalanced_dataset(seed=0) | ||
|
|
||
| baseline = AutoML() | ||
| baseline.fit(X_train=X, y_train=y, seed=42, **_fit_settings()) | ||
|
|
||
| resampled = AutoML() | ||
| resampled.fit( | ||
| X_train=X, | ||
| y_train=y, | ||
| resampler=SMOTE(random_state=42, k_neighbors=3), | ||
| seed=42, | ||
| **_fit_settings(), | ||
| ) | ||
|
|
||
| assert baseline.best_config != resampled.best_config, ( | ||
| "resampler=SMOTE(...) did not change the chosen best_config vs baseline; " | ||
| "the per-fold resampling hook may not be firing" | ||
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why are these changes needed?
Implements the design agreed on #1200 (option 3 — ship a minimal
resampler=kwarg, off by default). Closes the long-standing feature request frompetrosDemetrakopoulos(open since 2023) to plug an imbalanced-learn-style resampler into FLAML's cross-validation loop without leaking synthesized samples into the validation folds.The benchmark I posted earlier on two synthetic imbalance levels found that per-fold SMOTE doesn't materially outperform the pre-applied SMOTE workaround FLAML currently recommends, so this PR ships the integration as an opt-in convenience rather than a claimed quality improvement. Users who already roll their own per-fold pipelines get a cleaner path; users who don't get zero behavior change.
What the PR does
flaml/automl/automl.py—AutoML.fit:resampler=Nonekeyword-only argument (any object exposingfit_resample(X, y) -> (X, y)— matches imbalanced-learn'sBaseSamplerprotocol).ValueErrorif bothresamplerandsample_weightare passed (resampling breaks the 1-to-1 row alignment with weights; silent behavior in either direction would be a footgun).TypeErrorif the object doesn't exposefit_resample. Both errors fire atfit()entry, not mid-fold.task._resampler = resampler) so the per-fold hook can find it without a signature change downstream.flaml/automl/ml.py—get_val_loss:estimator.fit, if a resampler is set ontask, clones it viasklearn.base.clone(each fold gets an identical startingrandom_state; the caller-provided instance stays untouched) and callsX_train, y_train = fold_sampler.fit_resample(X_train, y_train).X_val, y_valare left at the raw class distribution — this is the entire point of the request and is asserted by thetest_resampler_leaves_validation_untouchedtest.evaluate_model_CV→get_val_loss) and the holdout path (compute_estimator→get_val_loss).test/automl/test_resampler.py— new file, five tests:test_resampler_changes_chosen_config— passingSMOTE(...)on a 6%-minority dataset picks a differentbest_configthan the same fit without the kwarg (proxy that the per-fold hook actually fires; a no-op would land on the identical config).test_resampler_leaves_validation_untouched— finalbest_loss(1 − F1) stays in[0, 1]on the raw imbalanced validation folds; a leak would drive it out of range or below the reachable minimum.test_resampler_with_sample_weight_raises— expectedValueError.test_resampler_without_fit_resample_raises— expectedTypeError.test_resampler_none_is_default_and_noop—resampler=Noneproduces identicalbest_config,best_estimator, andpredict(X)output as omitting the kwarg entirely; guarantees backward compatibility for existing users.imbalanced-learnis not added as a FLAML dependency. The tests usepytest.importorskipso they skip cleanly when it's not installed; end-users install it themselves if they want to pass a SMOTE object.Verified locally
pytest test/automl/test_resampler.py— 5/5 pass.pytest test/automl/test_split.py test/automl/test_preprocess_api.py— 18/18 pass.pre-commit run --files flaml/automl/automl.py flaml/automl/ml.py test/automl/test_resampler.py— all hooks pass.Related issue
Closes #1200.
Checks
#1200for the imbalanced-classification case; happy to expand that section in a follow-up now that the kwarg exists.)