feat: add optional resampler= kwarg to AutoML.fit for per-fold class-imbalance handling (#1200) by immu4989 · Pull Request #1568 · microsoft/FLAML

Imran Ahamed (immu4989) · 2026-07-02T00:03:19Z

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 from petrosDemetrakopoulos (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:

New resampler=None keyword-only argument (any object exposing fit_resample(X, y) -> (X, y) — matches imbalanced-learn's BaseSampler protocol).
Validation: raises ValueError if both resampler and sample_weight are passed (resampling breaks the 1-to-1 row alignment with weights; silent behavior in either direction would be a footgun).
Validation: raises TypeError if the object doesn't expose fit_resample. Both errors fire at fit() entry, not mid-fold.
Attaches the resampler to the task instance (task._resampler = resampler) so the per-fold hook can find it without a signature change downstream.

flaml/automl/ml.py — get_val_loss:

Just before estimator.fit, if a resampler is set on task, clones it via sklearn.base.clone (each fold gets an identical starting random_state; the caller-provided instance stays untouched) and calls X_train, y_train = fold_sampler.fit_resample(X_train, y_train).
X_val, y_val are left at the raw class distribution — this is the entire point of the request and is asserted by the test_resampler_leaves_validation_untouched test.
Single insertion point covers both the CV path (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 — passing SMOTE(...) on a 6%-minority dataset picks a different best_config than 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 — final best_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 — expected ValueError.
test_resampler_without_fit_resample_raises — expected TypeError.
test_resampler_none_is_default_and_noop — resampler=None produces identical best_config, best_estimator, and predict(X) output as omitting the kwarg entirely; guarantees backward compatibility for existing users.

imbalanced-learn is not added as a FLAML dependency. The tests use pytest.importorskip so they skip cleanly when it's not installed; end-users install it themselves if they want to pass a SMOTE object.

Verified locally

New tests: pytest test/automl/test_resampler.py — 5/5 pass.
Regression check: 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

I've used pre-commit to lint the changes in this PR (note the same in integrated in our CI checks).
I've included any doc changes needed for https://microsoft.github.io/FLAML/. See https://microsoft.github.io/FLAML/docs/Contribute#documentation to build and test documentation locally. (The Production-Deployment guide already points to #1200 for the imbalanced-classification case; happy to expand that section in a follow-up now that the kwarg exists.)
I've added tests (if relevant) corresponding to the changes introduced in this PR.
I've made sure all auto checks have passed.

…imbalance handling (microsoft#1200)

Copilot

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=None parameter to AutoML.fit, validate inputs, and store the resampler on the task for downstream access.
Apply per-fold resampling in get_val_loss by cloning the provided resampler and calling fit_resample(X_train, y_train) before estimator.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.

Imran Ahamed (immu4989) · 2026-08-23T12:34:13Z

        mlflow_logging=None,
        fit_kwargs_by_estimator=None,
        mlflow_exp_name=None,
+        resampler=None,
        **fit_kwargs,


Addressed in a40ebc0. The fit signature now places resampler after *, so it can only be passed as a keyword argument.

Imran Ahamed (immu4989) · 2026-08-23T12:34:17Z

+        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


Addressed in a40ebc0 and ab11996. The effective fit_kwargs_by_estimator settings are resolved before sample-weight validation, fit_resample must be callable, clone compatibility is checked at fit entry, and task._resampler is assigned on every fit so omitted values reset prior state.

Imran Ahamed (immu4989) · 2026-08-23T12:34:21Z

+    imblearn = pytest.importorskip("imblearn.over_sampling")
+    SMOTE = imblearn.SMOTE
+


Addressed in a40ebc0. The core tests now use a local deterministic DuplicateMinorityResampler, so they run without imbalanced-learn. A separate optional integration test still verifies SMOTE when imbalanced-learn is installed.

Imran Ahamed (immu4989) · 2026-08-23T12:34:24Z

+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"
+    )


Addressed in a40ebc0. The test now uses a custom metric to record each raw validation-fold class rate and asserts that it remains imbalanced, while the deterministic resampler counters verify that training partitions grow.

Imran Ahamed (immu4989) · 2026-08-23T12:34:27Z

+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"
+    )


Addressed in a40ebc0. The best-config comparison was removed and replaced with deterministic call-count and input/output-size assertions that directly verify each training-fold resampling hook.

…ble stub tests

…om/immu4989/FLAML into flaml-feature-1200-resampler-kwarg

Copilot

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

Suppressed comments (4)

flaml/automl/automl.py:2175

The user-facing Production-Deployment guide still says applying SMOTE upstream is the current recommendation and describes resampler= as a future integration (website/docs/Use-Cases/Production-Deployment.md:229). After this API ships, that guidance is stale and continues directing users toward the leakage-prone workaround; update the guide in this PR.

                by default. See issue #1200 for the design discussion and benchmarks.

flaml/automl/automl.py:2357

hasattr accepts a non-callable attribute, so a cloneable object with fit_resample = None passes entry validation and fails mid-fold with an unrelated 'NoneType' object is not callable error. Check callability to preserve the documented fail-fast behavior.

            if not hasattr(resampler, "fit_resample"):

flaml/automl/automl.py:2350

This validation checks only the argument passed directly to fit(), but fit_kwargs_by_estimator is resolved from AutoML(...) settings at line 2385. Thus AutoML(fit_kwargs_by_estimator={"lgbm": {"sample_weight": ...}}).fit(..., resampler=...) bypasses the error and reaches estimator fitting with weights shorter than the resampled data. Resolve the effective settings before checking them.

This issue also appears on line 2357 of the same file.

            weight_sources = [fit_kwargs] + list((fit_kwargs_by_estimator or {}).values())
            if any("sample_weight" in kw for kw in weight_sources):

flaml/automl/automl.py:2175

The public contract is inaccurate: a plain object exposing fit_resample is rejected unless sklearn.clone-compatible, while the implementation also applies the sampler to holdout evaluation rather than only CV folds. Document both constraints/paths so callers can select a compatible object and understand holdout behavior.

This issue also appears on line 2175 of the same file.

            resampler: object, default=None | An imbalanced-learn-compatible resampler
                (any object exposing `fit_resample(X, y) -> (X, y)`, such as
                `imblearn.over_sampling.SMOTE`). When set, the resampler is cloned and
                applied to each cross-validation fold's training partition before the
                estimator is fitted — validation partitions are left at the raw class

Imran Ahamed (immu4989) · 2026-08-23T12:34:31Z

+    resampler = getattr(task, "_resampler", None)
+    if resampler is not None:
+        from sklearn.base import clone
+
+        # Clone per fold/call so each estimator fit sees the same starting
+        # random_state; the caller-provided resampler stays untouched.
+        fold_sampler = clone(resampler)
+        X_train, y_train = fold_sampler.fit_resample(X_train, y_train)


Addressed in ab11996. A shared cloning helper now applies the resampler during fold evaluation, final and retrain fitting, and ensemble training. Counter-based regression tests cover the max_iter=1 final-fit path and ensemble training; all 13 focused tests and pre-commit checks pass.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

+                X_ensemble_train, y_ensemble_train = _resample_training_data(
+                    self._X_train_all,
+                    self._y_train_all,
+                    self._state.task,
+                )


 - `class_weight` is passed through to the underlying estimator unchanged if your chosen estimator accepts it (e.g., LightGBM, XGBoost sklearn API).

-For severe class imbalance, see also [issue #1200](https://github.com/microsoft/FLAML/issues/1200) on adding a `resampler=` integration. The current recommendation is to apply SMOTE (or your resampler of choice) upstream of `AutoML.fit`; see the imbalanced-learn documentation for the canonical pattern.
+For severe class imbalance, pass a cloneable imbalanced-learn sampler with `resampler=`:


Imran Ahamed (immu4989) and others added 2 commits July 1, 2026 19:02

feat: add optional resampler= kwarg to AutoML.fit for per-fold class-…

77adf31

…imbalance handling (microsoft#1200)

Merge branch 'main' into flaml-feature-1200-resampler-kwarg

4230e29

Li Jiang (thinkall) requested a review from Copilot July 16, 2026 01:45

Copilot started reviewing on behalf of Li Jiang (thinkall) July 16, 2026 01:47 View session

Copilot AI reviewed Jul 16, 2026

View reviewed changes

Li Jiang (thinkall) and others added 4 commits July 16, 2026 15:11

Merge branch 'main' into flaml-feature-1200-resampler-kwarg

80d10bd

address review: keyword-only resampler, stronger validation, CI-runna…

a40ebc0

…ble stub tests

Merge branch 'flaml-feature-1200-resampler-kwarg' of https://github.c…

950205d

…om/immu4989/FLAML into flaml-feature-1200-resampler-kwarg

Merge branch 'main' into flaml-feature-1200-resampler-kwarg

6ed1386

Li Jiang (thinkall) requested a balanced review from Copilot August 21, 2026 07:03

Copilot started reviewing on behalf of Li Jiang (thinkall) August 21, 2026 07:07 View session

Copilot AI reviewed Aug 21, 2026

View reviewed changes

fix: apply resampler to final model training

ab11996

Li Jiang (thinkall) requested a balanced review from Copilot August 24, 2026 02:26

Copilot started reviewing on behalf of Li Jiang (thinkall) August 24, 2026 02:27 View session

Copilot AI reviewed Aug 24, 2026

View reviewed changes

Navigation Menu

Uh oh!

feat: add optional resampler= kwarg to AutoML.fit for per-fold class-imbalance handling (#1200) - #1568

feat: add optional resampler= kwarg to AutoML.fit for per-fold class-imbalance handling (#1200)#1568
Imran Ahamed (immu4989) wants to merge 7 commits into
microsoft:mainfrom
immu4989:flaml-feature-1200-resampler-kwarg

Imran Ahamed (immu4989) commented Jul 2, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		imblearn = pytest.importorskip("imblearn.over_sampling")
		SMOTE = imblearn.SMOTE

Uh oh!

Conversation

Imran Ahamed (immu4989) commented Jul 2, 2026

Why are these changes needed?

What the PR does

Verified locally

Related issue

Checks

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Imran Ahamed (immu4989) Aug 23, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants