What’s Taija?

Taija is a small but growing ecosystem of packages geared towards Trustworthy AI in JuliA—Taija

Trustworthy AI in Julia: github.com/JuliaTrustworthyAI

Trustworthy AI in Julia: github.com/JuliaTrustworthyAI

Ecosystem Overview

%%{
  init: {
    'theme': 'base',
    'themeVariables': {
      'primaryColor': '#BB2528',
      'primaryTextColor': '#fff',
      'primaryBorderColor': '#7C0000',
      'lineColor': '#F8B229',
      'secondaryColor': '#006100',
      'tertiaryColor': '#e9edfb',
      'fontFamily': "avenir"
    }
  }
}%%

flowchart TB

    classDef taija fill:#389836,stroke:#333,color:#fff;
    classDef core fill:#CB3C33,stroke:#333,color:#fff;
    classDef base fill:#9558B2,stroke:#333,color:#fff;

    %% Base
    base["TaijaBase.jl"]

    %% Meta
    interop["TaijaInteroperability.jl"]
    data["TaijaData.jl"]
    parallel["TaijaParallel.jl"]
    plotting["TaijaPlotting.jl"]

    %% Core
    ce["CounterfactualExplanations.jl"]
    cp["ConformalPrediction.jl"]
    lr["LaplaceRedux.jl"]
    jem["JointEnergyModels.jl"]

    %% External
    mlj["MLJ.jl"]
    flux["Flux.jl"]

    class base base;
    class interop,data,parallel,plotting taija;
    class ce,cp,lr,jem core;

    %% Graph
    base --> ce & cp & lr & jem

    subgraph "Core Packages"
        ce & cp & lr & jem 
    end

    subgraph "Meta Packages"
        data & plotting & parallel & interop
    end 

    subgraph "External Packages"
        mlj & flux
    end

Overview of the Taija ecosystem. Early-stage packages omitted.

Ecosystem Overview

%%{
  init: {
    'theme': 'base',
    'themeVariables': {
      'primaryColor': '#BB2528',
      'primaryTextColor': '#fff',
      'primaryBorderColor': '#7C0000',
      'lineColor': '#F8B229',
      'secondaryColor': '#006100',
      'tertiaryColor': '#e9edfb',
      'fontFamily': "avenir"
    }
  }
}%%

flowchart TB

    classDef taija fill:#389836,stroke:#333,color:#fff;
    classDef core fill:#CB3C33,stroke:#333,color:#fff;
    classDef base fill:#9558B2,stroke:#333,color:#fff;

    %% Base
    base["TaijaBase.jl"]

    %% Meta
    interop["TaijaInteroperability.jl"]
    data["TaijaData.jl"]
    parallel["TaijaParallel.jl"]
    plotting["TaijaPlotting.jl"]

    %% Core
    ce["CounterfactualExplanations.jl"]
    cp["ConformalPrediction.jl"]
    lr["LaplaceRedux.jl"]
    jem["JointEnergyModels.jl"]

    %% External
    mlj["MLJ.jl"]
    flux["Flux.jl"]

    class base base;
    class interop,data,parallel,plotting taija;
    class ce,cp,lr,jem core;

    %% Graph
    base --> ce & cp & lr & jem

    subgraph "Core Packages"
        ce & cp & lr & jem 
    end

    subgraph "Meta Packages"
        data & plotting & parallel & interop
    end 

    subgraph "External Packages"
        mlj & flux
    end

Overview of the Taija ecosystem. Early-stage packages omitted.

Who’s behind Taija?

  • 2021-2022: Initially developed by myself to support my PhD research in Trustworthy AI (still ongoing).
  • 2023-2024: Expanded by a growing community of contributors, including TU Delft and G/JSoC students.

Thanks to @MojiFarmanbar, @JorgeLuizFranco, @Rockdeldiablo, @kmariuszk, @RaunoArike, @VincentPikand, @severinbratus, @rithik83, @navimakarov, @laurikskl, @MarkArdman, @adelinacazacu, @Andrei32Ionescu and many others!

Use Cases

Who could benefit from Taija?

  • Researchers in AI and ML, particularly in the fields of explainability, uncertainty quantification, and Bayesian deep learning: (, …)
  • Practitioners using conventional ML and DL models who are interested in understanding the models’ decisions and their uncertainty.
  • Julia developers who want to contribute to the ecosystem (any level of expertise is welcome!).

Research

Counterfactual Explanations

The largest single category of CE methods solves the following optimization through gradient descent:

\[ \begin{aligned} \mathbf{s}^* &= \arg \min_{\mathbf{s}^\prime \in \mathcal{S}} \left\{ {\text{yloss}(M(f(\mathbf{s}^\prime)),y^*)}+ \lambda {\text{cost}(f(\mathbf{s}^\prime)) } \right\} \end{aligned} \]

Counterfactual Explanations explain how inputs into a model need to change for it to produce different outputs.

Figure 1: Gradient-based counterfactual search.

Pick your Poison

All of these counterfactuals are valid explanations …

… which one would you pick?

Figure 2: Turning a 9 into a 7: Counterfactual explanations for an image classifier produced using Wachter (Wachter, Mittelstadt, and Russell 2017), Schut (Schut et al. 2021) and REVISE (Joshi et al. 2019).

Faithful Counterfactuals

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

ECCCo1 counterfactuals

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

Intent Classification

Intent classification (IC) in dialogue systems is a common task and a natural place for conformal prediction.

  • Simply returning top-1 softmax likely wrong.
  • Existing ad-hoc approach is top-\(k\).
  • Conformal classifiers predict sets that fulfill coverage guarantee.
Figure 5: A simple conformal chat bot in the Julia REPL using ConformalPrediction.jl.

Conformal IC and Clarification

  • Our NAACL 2024 paper introduces CICC: a framework for fast and accurate intent classification in conversational AI.
  • Winning project at ING Experiment Week 2023.

See also Building a Conformal Chatbot in Julia with ConformalPrediction.jl and Transformers.jl1.

More Research

Developments

A brief overview of the highlights. For a more introductory presentation, see also this slides presentation The Alan Turing Institute

What’s new in …

CounterfactualExplanations.jl: A package for Counterfactual Explanations and Algorithmic Recourse in Julia.

  • refactoring (e.g. package extensions)
  • performance improvements
  • new features

Composable Generators

Recall that for most generators, we have:

\[ \begin{aligned} \mathbf{s}^* &= \arg \min_{\mathbf{s}^\prime \in \mathcal{S}} \left\{ {\text{yloss}(M(f(\mathbf{s}^\prime)),y^*)}+ \lambda {\text{cost}(f(\mathbf{s}^\prime)) } \right\} \end{aligned} \]

Why not compose generators that combine ideas from different off-the-shelf generators?

@chain generator begin
    @objective logitcrossentropy 
      + 1.0ddp_diversity     # DiCE (Mothilal et al. 2020)
    @with_optimiser Flux.Adam(0.1)                      
    @search_latent_space     # REVISE (Joshi et al. 2019)
end

Explaining Different Models

Besides any Flux.jl model, extensions add support for

Figure 6: Counterfactual for differentiable decision tree classifier.

Benchmarking Explanations

Extensive support for evaluating and benchmarking explanations.

Evaluation

# Generate counterfactuals
ces = generate_counterfactual(
  factual, 
  target_label, 
  data, 
  M, 
  generator; 
  num_counterfactuals=5
)

# Evaluate them
evaluate(ces)

Benchmarks

Benchmark all available generators and models at once in parallel1:

using TaijaParallel
pllr = ThreadsParallelizer()
bmk = benchmark(
  counterfactual_data; 
  parallelizer = pllr
)

What’s new in …

LaplaceRedux.jl: A package for Effortless Bayesian Deep Learning through Laplace Approximation for Flux.jl neural networks.

  • new features
  • interface to MLJ
  • JCon Proceedings (under review).

Student Contributions

  • Support for multi-class problems.
  • Support for more scalable Hessian approximations.
  • Interface to MLJ for easy model training and evaluation.

Figure 7: Hessian approximations. Source: Daxberger et al. (2021)

Check out their blog post!

What’s new in …

ConformalPrediction.jl: A package for Predictive Uncertainty Quantification through Conformal Prediction for Machine Learning models trained in MLJ.

Joint Energy Models

JointEnergyModels.jl: A package for Joint Energy-Based Models Models in Julia (early development).

Hybrid models that can predict and generate (Grathwohl et al. 2020).

  • Flux.jl interface
  • MLJFlux.jl interface
Figure 8: Predicted output class probabilities (contour) and generated inputs (stars) for a binary JEM classifier.

Final Things

Ongoing Work

Taija has been running two Julia Season of Code projects this summer.

  1. (Conformal Bayes) Bridging the gap between Bayesian and frequentist approaches to Predictive Uncertainty Quantification with @Rockdeldiablo and co-supervisor @MojiFarmanbar
  2. (Causal Recourse) From minimal perturbations to minimal interventions for Algorithmic Recourse with @JorgeLuizFranco and co-supervisor @mschauer.

Student Testimonials

Students have generally been enthusiastic about their experience with Julia and Taija:

“Programming in Julia has definitely helped us become better programmers. […] whenever we had such questions and asked them [to] the wider Julia community, there were always people ready to help in my experience, which was nice.”

— @RaunoArike

Get Involved

  • Working on related projects?
  • Interested in contributing to Taija?
  • Want to learn more about Trustworthy AI in Julia?
  • Any suggestions or feedback?

Get in touch with me or any of the contributors! Join our #taija channel on the JuliaLang Slack or visit our GitHub organization.

Questions?

References

Altmeyer, Patrick, Giovan Angela, Aleksander Buszydlik, Karol Dobiczek, Arie van Deursen, and Cynthia CS Liem. 2023. “Endogenous Macrodynamics in Algorithmic Recourse.” In 2023 IEEE Conference on Secure and Trustworthy Machine Learning (SaTML), 418–31. IEEE.
Altmeyer, Patrick, Andrew M. Demetriou, Antony Bartlett, and Cynthia C. S. Liem. 2024. “Position Paper: Against Spurious Sparks-Dovelating Inflated AI Claims.” https://arxiv.org/abs/2402.03962.
Altmeyer, Patrick, Arie van Deursen, et al. 2023. “Explaining Black-Box Models Through Counterfactuals.” In Proceedings of the JuliaCon Conferences, 1:130. 1.
Daxberger, Erik, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. 2021. “Laplace Redux-Effortless Bayesian Deep Learning.” Advances in Neural Information Processing Systems 34.
Grathwohl, Will, Kuan-Chieh Wang, Joern-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. 2020. “Your Classifier Is Secretly an Energy Based Model and You Should Treat It Like One.” In International Conference on Learning Representations.
Hengst, Floris den, Ralf Wolter, Patrick Altmeyer, and Arda Kaygan. 2024. “Conformal Intent Classification and Clarification for Fast and Accurate Intent Recognition.” https://arxiv.org/abs/2403.18973.
Joshi, Shalmali, Oluwasanmi Koyejo, Warut Vijitbenjaronk, Been Kim, and Joydeep Ghosh. 2019. “Towards Realistic Individual Recourse and Actionable Explanations in Black-Box Decision Making Systems.” https://arxiv.org/abs/1907.09615.
Schut, Lisa, Oscar Key, Rory Mc Grath, Luca Costabello, Bogdan Sacaleanu, Yarin Gal, et al. 2021. “Generating Interpretable Counterfactual Explanations By Implicit Minimisation of Epistemic and Aleatoric Uncertainties.” In International Conference on Artificial Intelligence and Statistics, 1756–64. PMLR.
Wachter, Sandra, Brent Mittelstadt, and Chris Russell. 2017. “Counterfactual Explanations Without Opening the Black Box: Automated Decisions and the GDPR.” Harv. JL & Tech. 31: 841. https://doi.org/10.2139/ssrn.3063289.
Zagorac, Ivor. 2024. “A Study on Counterfactual Explanations: Investigating the Impact of Inter-Class Distance and Data Imbalance.”