Generalize DiCE-Torch to Support Multiclass Probability Targets

[sofeikov]

Hi all,

I’m the author of #454.
It’s been a while since our last discussion, but I’d like to revisit this and get your input before I start contributing code.

Context

I was experimenting with DiCE in a heterogeneous graph learning setting, where I wanted to see which node features need to change to achieve a desired probability. My problem involved 3 classes.

Current limitations I ran into

The dice-torch class appears to only support binary classification, treating one outcome as “bad” (prob→0) and the other as “good” (prob→1). This prevents straightforward use in multiclass problems.

Arbitrary thresholds passed by the user are ignored and defaulted to 0.75.

The current setup makes it hard to:

• target changes in probabilities (e.g., increase a class probability by +7%),
• or return “best-effort” counterfactuals when thresholds are unattainable.

Proposal

• Generalize the treatment of outcomes from a scalar probability to a vector-based format.
• Most classifiers already output softmax vectors.
• If we instead optimize a target index of the probability vector (provided at construction), we can support multiclass problems directly.

Add options to:

• target examples that gave the largest probability for the desired class (vs. closest to threshold),
• store and return best-effort counterfactuals when thresholds cannot be met,
• support arbitrary probability thresholds (e.g., “increase probability of class X by 7%”).

Benefits

• Extends DiCE beyond binary-only settings into general multiclass classification.
• Enables use in probabilistic workflows (e.g., Monte Carlo simulations), where shifts in probability distributions are more meaningful than passing a single threshold.
• Reuses most of DiCE’s existing machinery with minimal conceptual changes.

If this direction makes sense to you, I’d be happy to start drafting test cases and working on a prototype implementation.

Thanks,
Konstantin

[amit-sharma]

@sofeikov thanks for starting this thread. Agree with the proposal here. dice-torch only supports binary classification, and it will be great to extend to it multi-class.

Just one suggestion: as you are making changes, try to keep the user-facing API almost the same for a user using binary classification. So that this update does not include a breaking change for people using the current DiCE version.

[sofeikov]

Fair enough, I will start with something really simple, like keeping the user-passed thresholds respected, we align and go form there 👍

[sofeikov]

Hi there,

It's been a while here :)

Quick update from my side: since opening this, I ended up implementing a good part of this direction in my fork because I needed it for my own workflows.

On the counterfactual/API side, the work I have locally now includes:

• respecting user-passed thresholds instead of overriding them,
• multiclass support built around explicit target classes / label-indexed probability targets,
• relative probability uplift targets (for example, asking for a +7% increase),
• an explicit best-effort mode when the requested threshold is unreachable,
• a selection mode that can maximize desired-class probability instead of just picking examples closest to the threshold,
• per-feature direction constraints such as increase-only / decrease-only.

Separately, in my fork I also migrated the project to uv, modernized the pandas-related parts of the repo, and removed TensorFlow support to keep the codebase aligned with the workflows I actually use.

To be clear, I did this primarily for myself so I could keep moving with my own use cases, not because I assumed all of it should land upstream as one big change.

If any of the above is remotely interesting for you, can I suggest to take a look at the diff between the two repos(I tried to keep the formatting similar so the diff is easier to review and it might be easiest to start with the README diff) and if any of that work looks relevant to the repo, I'm happy to submit clean, isolated PRs for the upstream.

Let me know if anything :)

Fork for reference: https://github.com/sofeikov/DiCE

[amit-sharma]

Hey @sofeikov thanks for sharing this. Your improvements are very useful, and I believe many others would also benefit. I'll take a detailed look at the code and share suggestions.