GPU compatible (L-)C2ST

[JuliaLinhart]

Is your feature request related to a problem? Please describe.

When working with high-dimensional datasets, computing C2ST (and particularly L-C2ST) metrics can be quite slow. This is primarily due to the reliance on scikit-learn (sklearn) classifiers, which do not support GPU acceleration. As a result, performance is limited, especially for large datasets or complex models.

Describe the solution you'd like

Introducing support for classifiers from GPU-compatible libraries, such as skorch, would significantly improve performance. By leveraging GPU acceleration, training and evaluation of classifiers for C2ST metrics could become substantially faster, enhancing both development and experimentation workflows.

Describe alternatives you've considered

Currently, no viable alternatives are available since the limitation stems from the lack of GPU support in sklearn. Expanding support to include GPU-compatible libraries appears to be the most effective solution.

Additional context

• [ ] Would require adding another library as a dependency. Need to check if this might conflict with our current dependencies.
• [ ] Adapt current implementation here
• [ ] Think about similar accelerate similar metrics i.e. here

[Dev-Sudarshan]

Hey @JuliaLinhart , I would love to try to solve this issue .While exploring ways to implement GPU support via skorch, I noticed that skorch only wraps PyTorch neural networks and cannot replace tree-based classifiers. Since C2ST uses Random Forests but L-C2ST uses MLPs, skorch can be applied to L-C2ST but not directly to C2ST.

Could you please advise on how you would like me to proceed?

[JuliaLinhart]

Hello! That's amazing, thanks for the initiative!

1. I’d apply skorch to L-C2ST or MLP-based C2ST to enable GPU acceleration

2. For RF-based C2ST, we have a few options:

• Keep using CPU-based Random Forests from sklearn and consider Intel acceleration scikit-learn-intelex
• Try a GPU-compatible alternative, such as RAPIDS cuML RandomForestClassifier or XGBoost with gpu_hist, both offering significant speedups on large datasets.

In any case we have to think about how adding those changes would affect the API (conflict with current dependencies, userfriendly, etc...). We would have to check with @janfb what he thinks would be best.

[janfb]

Great that you tackle this @Dev-Sudarshan !

I agree with @JuliaLinhart , we can just focus on MLP-based C2ST for the GPU support for L-C2ST.

The different C2ST options are independent of L-C2ST in principle, and we could just issue a user warning when someone tries to use L-C2ST with RF-based C2ST on the GPU.

So, I would recommend focusing on using skorch and supporting only the MLP-based C2ST, and keep the API changes minimal, e.g., just a flag which device and careful device handling similar to how we do it in the inference methods.

@JuliaLinhart would you be up for guiding this? and I can do a high-level review later?

[JuliaLinhart]

Yes sounds good!!

[Dev-Sudarshan]

I'm fairly new to this topic, so I initially sought help from chat-based models to understand the structure and requirements. After reviewing their suggestions and taking time to understand the logic myself, I believe the following changes are needed to enhance lc2st.py—especially for integrating GPU support and PyTorch compatibility. @JuliaLinhart Could you please review these before I move forward with implementation?

1. Change in Code Dependencies Imports: Add torch, torch.nn, and skorch.NeuralNetClassifier. Utilities: Define a simple get_device() function to determine if CUDA is available. Model: Define the PyTorchMLP class (nn.Module)—the actual network that will run on the GPU.

2. LC2ST.init Logic Add device: str = 'cpu' to init parameters. Conditional Logic (for 'mlp'): IF device == 'cuda': Set self.clf_class to NeuralNetClassifier and configure self.clf_kwargs with PyTorch-specific params (module=PyTorchMLP, device='cuda', etc.). ELSE (CPU Default): Set self.clf_class to MLPClassifier (standard sklearn).

3. train_lc2st Function Data Type Fix: Change label creation to explicitly use dtype=np.int64 (long integers). This is required for PyTorch/skorch classification.

[JuliaLinhart]

Yes, so if I understand correctly (I havn't used skorch before either...) you need a PyTorch nn.Module to wrap the skorch.NeuralNetClassifier around, like this as mentioned here:

class MyModule(nn.Module):
    def __init__(self, num_units=10, nonlin=nn.ReLU()):
        super().__init__()

        self.dense0 = nn.Linear(20, num_units)
        self.nonlin = nonlin
        self.dropout = nn.Dropout(0.5)
        self.dense1 = nn.Linear(num_units, num_units)
        self.output = nn.Linear(num_units, 2)

    def forward(self, X, **kwargs):
        X = self.nonlin(self.dense0(X))
        X = self.dropout(X)
        X = self.nonlin(self.dense1(X))
        X = self.output(X)
        return X


net = NeuralNetClassifier(
    MyModule,
    max_epochs=10,
    criterion=nn.CrossEntropyLoss(),
    lr=0.1,
    # Shuffle training data on each epoch
    iterator_train__shuffle=True,
)

Then I guess all you are saying should be true.

I suggest you try it on your end: fork the repo, implement your changes, run pytests. On that note, you might need to change the existing pytests and/or add some new ones to check for device compatibility etc. (cf what janbf mentioned related to what we do in the inference methods).

Also try to think of appropriate user warnings like this one janbf mentioned: "The different C2ST options are independent of L-C2ST in principle, and we could just issue a user warning when someone tries to use L-C2ST with RF-based C2ST on the GPU."

[Dev-Sudarshan]

Hi @JuliaLinhart , I’ve been implementing the GPU support for L-C2ST. I noticed a potential parameter handling issue in LC2ST.init:

The default classifier_kwargs (for the CPU MLPClassifier), including the scaled hidden_layer_sizes and training settings, are currently only set if classifier_kwargs is None.

If a user provides custom classifier_kwargs (e.g., to override one parameter like max_iter), all other essential sbi default parameters are skipped. This forces the classifier to use standard scikit-learn defaults, which are different and potentially less suitable for SBI (e.g., smaller fixed network size, fewer training iterations).

Before I add logic to merge the user's classifier_kwargs with our SBI defaults, could you confirm if this parameter merging is already being handled in a higher-level wrapper function that calls LC2ST?

[JuliaLinhart]

Yes, that's normal. I wanted the code to be as flexible and transparent as possible:

• for those who have no intuition on which kwargs to use, defaults to the sbi ones
• for those who know what they want: include the kwargs themselves

So know mixture of both, where one might miss hidden kwarg choices. Makes sense?

It's a personal choice though, might want to discuss this with @janfb .

[janfb]

I suggest we find an implementation that uses the sbi settings by default, and when a user passes additional options, those settings will be updated for the kwargs the user passed. The other kwargs should not be overwritten but used as defined in the sbi (potentially different from the scikit-learn default).

This should be possible using dict.update, or is it more involved than that?

[Dev-Sudarshan]

To put the merging logic in, I recommend using the dictionary merge operator ({**D1, **D2}) instead of the dict.update() function.

Using update() can accidentally change the default settings permanently. If one user runs the code and changes a default, the next user (or the next iteration) will get the wrong setting.

The merge operator is safer because it leaves the original defaults untouched. Does that sound okay for moving forward?

[janfb]

I had a closer look again and want to summarize the problem again, just for the record.

Julia’s original idea (either full defaults or full control) is reasonable, but for most users it will be more intuitive if passing a single kwarg (like max_iter) keeps the other sbi-tuned settings instead of falling back to plain scikit-learn defaults.

On the dict.update vs {**D1, **D2} question: the important part is that we construct a fresh kwargs dict per LC2ST instance and don’t mutate any shared global defaults. If we do that, both styles are fine. So feel free to use the merge operator if you find it clearer.

I’d suggest you go ahead and implement the merging logic along those lines (plus a small test that checks “defaults only” and “defaults + single override”), and we can discuss any finer details in the PR. OK?

[JuliaLinhart]

Sounds all good to me. Just be careful about one thing: if the user wants to change the default kwargs he should be allowed to overwrite them. Without changing it permanently and impacting the next user ofc!