Counterfactual Training

Teaching Models Plausible and Actionable Explanations using Taija’s new package CounterfactualTraining.jl

Patrick Altmeyer

Delft University of Technology

2026-03-23

Taija

Taija hosts research software geared towards Trustworthy Artificial Intelligence in Julia. Developed during my PhD.

• Model Explainability (CounterfactualExplanations.jl)
• Predictive Uncertainty Quantification (ConformalPrediction.jl, LaplaceRedux.jl)
• Effortless Bayesian Deep Learning (LaplaceRedux.jl)
• Hybrid Learning (JointEnergyModels.jl, CounterfactualTraining.jl)

From Counterfactual Explanations to Counterfactual Training

Training Opaque Models

Tweaking Parameters

Objective:

\[ \begin{aligned} \min_{\textcolor{orange}{\theta}} \{ {\text{yloss}(M_{\theta}(\mathbf{x}),\mathbf{y})} \} \end{aligned} \]

Training Opaque Models

Tweaking Parameters ::: {.columns} ::: {.column width=75%} Objective:

\[ \begin{aligned} \min_{\textcolor{orange}{\theta}} \{ {\text{yloss}(M_{\theta}(\mathbf{x}),\mathbf{y})} \} \end{aligned} \]

Solution:

\[ \begin{aligned} \theta_{t+1} &= \theta_t - \nabla_{\textcolor{orange}{\theta}} \{ {\text{yloss}(M_{\theta}(\mathbf{x}),\mathbf{y})} \} \\ \textcolor{orange}{\theta^*}&=\theta_T \end{aligned} \] ::: ::: {.column width=25%}

::: :::

Explaining Opaque Models

Tweaking Inputs–CounterfactualExplanations.jl

Objective:

\[ \begin{aligned} \min_{\textcolor{purple}{\mathbf{x}}} \{ {\text{yloss}(M_{\textcolor{orange}{\theta^*}}(\mathbf{x}),\mathbf{y^{\textcolor{purple}{+}} }) + \lambda \text{reg}(\mathbf{x};\cdot) } \} \end{aligned} \]

Explaining Opaque Models

Tweaking Inputs–CounterfactualExplanations.jl

Objective:

\[ \begin{aligned} \min_{\textcolor{purple}{\mathbf{x}}} \{ {\text{yloss}(M_{\textcolor{orange}{\theta^*}}(\mathbf{x}),\mathbf{y^{\textcolor{purple}{+}} }) + \lambda \text{reg}(\mathbf{x};\cdot)} \} \end{aligned} \]

Solution:

\[ \begin{aligned} \mathbf{x}_{t+1} &= \mathbf{x}_t - \nabla_{\textcolor{purple}{\mathbf{x}}} \{ \text{yloss}(M_{\textcolor{orange}{\theta^*}}(\mathbf{x}),\mathbf{y^{\textcolor{purple}{+}} }) \\&+ \lambda \text{reg}(\mathbf{x};\cdot) \} \\ \textcolor{purple}{\mathbf{x}^*}&=\mathbf{x}_T \end{aligned} \]

Explanation or Adversarial Example?

Plausibility at all cost?

All of these counterfactuals are valid explanations for the model’s prediction.

Pick your poison …

Figure 1: Turning a 9 into a 7: Counterfactual explanations for an image classifier using different approaches (Altmeyer et al. 2024).

Faithful First, Plausible Second

Figure 2: Turning a 9 into a 7. ECCCo applied to MLP (a), Ensemble (b), JEM (c), JEM Ensemble (d).

Insight: faithfulness facilitates¹

• model quality checks (Figure 2).
• state-of-the-art plausibility (Figure 3).

Figure 3: Results for different generators (from 3 to 5).

1. Altmeyer et al. (2024) @ AAAI 2024. [blog]; ECCoGenerator in CounterfactualExplanations.jl

Putting it all together

Counterfactual Training

First, Tweaking Inputs¹

\[ \begin{aligned} \mathbf{x}_{t+1} &= \mathbf{x}_t - \nabla_{\textcolor{purple}{\mathbf{x}}} \{ {ECCCo(M_{\textcolor{orange}{\theta^*}}(\mathbf{x}),\mathbf{y^{\textcolor{purple}{+}} })} \} \\ \textcolor{purple}{\mathbf{x}^*}&=\mathbf{x}_T \end{aligned} \]

Then, Tweaking Parameters

\[ \begin{aligned} \theta_{t+1} &= \theta_t - \nabla_{\textcolor{orange}{\theta}} \{ {\text{yloss}(M_{\theta}(\mathbf{x}),\mathbf{y})} + \text{div}(\textcolor{purple}{\mathbf{x}^*},\mathbf{x}^+,y^+; \theta) \} \\ \textcolor{orange}{\theta^*}&=\theta_T \end{aligned} \]

1. Generate faithful explanations using ECCCo objective (Altmeyer et al. 2024).

Counterfactual Training

1. Contrast faithful CE with data \(\rightarrow\) Explainability \(\uparrow\)
2. Feature mutability constraints \(\rightarrow\) Actionability \(\uparrow\)(holds provably under certain assumptions)
3. Bonus: use nascent CE as AE \(\rightarrow\) Robustness \(\uparrow\)

Figure 4: (a) conventional training, all mutable; (b) CT, all mutable; (c) conventional, age immutable; (d) CT, age immutable. CE generated using ECCoGenerator.

Counterfactual Training: Results

Plausibility: Visual explanations (counterfactuals) for baseline (top row) vs CT (bottom).

Actionability: Visual explanations (integrated gradients) as before. Five top and bottom rows immutable.

Test accuracies on adversarially perturbed data with varying perturbation sizes.

The Julia package

Main API

CounterfactualTraining.jl can be used with deep learning models from the FluxML ecosystem.

using CounterfactualTraining
using CounterfactualExplanations
using Flux

model, log = counterfactual_training(
  obj,        # custom type
  model,      # Flux.jl model
  generator,  # CE.jl type
  train_set,  # MLUtils.jl dataloader
  opt_state,  # Optimisers.jl
)

Research Branch

• Defining feature: many small models + many counterfactuals + different generators tested
• Solution: CounterfactualExplanations.jl (JuliaCon 2022) + TaijaParallel.jl (JuliaCon 2024) support multi-processing through MPI.jl.

Parallel CPU loop

for epoch in start:end
  # 1. Generate counterfactuals in parallel across CPUs
  for (i, batch) in enumerate(train_set)
    # 2. Unwrap and backprop
  end
end

Research Branch

Pros	Cons
✅ Access to comprehensive CE.jl generator suite	❌ Bottlenecked by CE.jl performance
✅ Viable alternative for GPU-poor	❌ Large models
✅ Generate CEs in parallel	❌ Generate CEs on GPU

GPU Branch

(a) In assembly.

(b) In its full glory.

Figure 5: Entering a brave new GPU-not-so-poor world with the Framework Desktop.

GPU Branch

• Adds new NativeGenerator: GPU-friendly re-implementation of CE.jl’s ECCoGenerator.
• Overloads counterfactual_training to (1) generate counterfactuals on GPU, (2) train model on GPU.

generator = NativeGenerator()

model, log = counterfactual_training(
  obj,        # custom type
  model,      # Flux.jl model
  generator,  # custom type 
  train_set,  # MLUtils.jl dataloader
  opt_state,  # Optimisers.jl
)

GPU Branch

ResNet for MNIST anyone? New branch supports large neural networks (even when you don’t need them).

Figure 6: Per-epoch wall-clock time, training accuracy, and implausibility trajectory for the full and vanilla objectives. Complete tutorial available in the docs.

Contributions Welcome!

Taija was my PhD project. Unfortunately, I am more constrained these days 😔

Performance-focused

1. Improve performance of CE.jl #526 (CE.jl)
2. Make CE.jl GPU-friendly #262 (CE.jl)

Research-focused

3. Per-epoch vs. per-batch CE generation #107 (CT.jl)
4. Doing related XAI research? #103 (CE.jl)

Check it out!

Preprint

Software

Homepage

References

Altmeyer, Patrick, Mojtaba Farmanbar, Arie van Deursen, and Cynthia C. S. Liem. 2024. “Faithful Model Explanations through Energy-Constrained Conformal Counterfactuals.” Proceedings of the Thirty-Eighth AAAI Conference on Artificial Intelligence 38: 10829–37. https://doi.org/10.1609/aaai.v38i10.28956.