Numerical stability of `confusion_matrix_at_thresholds` on `float32`-only capable array API devices

[ogrisel]

Warning

This issue is not yet ready for a PR. If you are interested in contributing to scikit-learn, please have a look at our contributing guidelines, and in particular the sections for new contributors and the "Needs triage" label.

As discussed in #33200, confusion_matrix_at_thresholds currently relies on an upcast to _max_precision_float_dtype(xp, device) which is float64 for numpy on CPU and many array API devices for can be limitted to float32 on common devices such as torch MPS (Apple GPUs) or some older Intel GPUs via torch XPI.

Possible solution

I think this code could be changed to improve numerical stability (and maybe processing speed) by implementing either or both of the following two strategies:

When sample_weight is None

Using an integer dtype to perform the call to xp.cumulative_sum instead and only after casting the results to the dtype of y_score.

When sample_weight is not None

Remove the upcast to _max_precision_float_dtype(xp, device) and instead compute the total weight sum and normalize the weight array before calling the cumulative sum. Here is the pseudo-code tps = xp.cumulative_sum(y_true * weight / weigh.sum(), dtype=y_score.dtype) * weight.sum().

Resolution plan

I think the first step is to do a quick empirical study on a float32 only device to confirm that this function actually has a numerical stability problem. Instead of MPS, we could leverage array-api-strict's float32-only devices and call that function on arrays with a large number of elements to try to trigger the problem and then check that the proposed solution actually work as intended. If this is the case we could open two independent PRs, one for the weighted case and one for the unweighted case.

[Kaynup]

Hello @ogrisel, I ran an empirical investigation on the unweighted case (sample_weight=None) to verify the numerical stability issue on float32 devices and evaluate the proposed integer accumulation strategy.

I ran a benchmark on 20,000,000 samples and tested for both cases:

# Float32 accumulation
tps_float32 = np.cumsum(y_true.astype(np.float32), dtype=np.float32)[-1]
# 16,777,216.0 (limits at 2^24, losing rest 3,222,784 samples)

# Integer accumulation (Proposed strategy)
tps_int64 = np.cumsum(y_true.astype(np.int64), dtype=np.int64)[-1]
# 20,000,000 (0% loss)

I've opened Draft PR #34817 only for reference and share the implementation only.
I validated some edge cases along with it like, threshold grouping slicing for integer cum-sum, explicit casting to avoid mixed subtraction errors in array-api-strict, and casting to y_score.dtype to preserve output precision.

[ogrisel]

Once #34817 is merged to tackle the unweighted case, we can rebase #34827 to tackle the weighted case with a slightly more complex strategy that leverages a fixed point integer scaling strategy.

The weight normalization method I suggested originally in the description does not seem to yield a robust enough fix for this problem.