Google Summer of Code 2024 final report

GSoC2024finalreport.md

Google Summer of Code 2024 Final Report

• Name: Pasquale Caterino
• Organization: Taija (Trustworthy Artificial Intelligence in Julia)
• Mentors: Patrick Altmeyer and Mojtaba Farmanbar
• Project: [Add support for Conformalized Bayes to ConformalPrediction.jl](Add support for Conformalized Bayes to ConformalPrediction.jl
• Github general repo :(https://github.com/JuliaTrustworthyAI)

Project Description

This project took place over the summer of 2024 as part of the Google Summer of Code initiative, with the support of the Taija (Trustworthy Artificial Intelligence in Julia) initiative. Taija is an organization whose aim is to provide users with the tools they need to understand opaque machine learning models and interpret their results so that thoughtful decisions can be made based on a probabilistic interpretation of the results. For this purpose, the organization provides several packages for various tasks, like counterfactual explanations (CounterfactualExplanations.jl), approximate Bayesian inference through the Laplace approximation (LaplaceRedux.jl), Algorithmic Recourse (AlgorithmicRecourseDynamics.jl) and conformal predictions (ConformalPrediction.jl). Alongside these packages there are other support packages created to extend the functionalities of the core packages.

The original goal of the project was to improve and expand the capabilities of the package ConformalPrediction.jl to support conformalized Bayes. This technique makes it possible to combine the conformal prediction technique, which is based on a frequentist interpretation of statistics, with Bayesian inference theory to obtain conformal sets whose width depends on the posterior predictive distribution obtained through the Bayesian approach. For this purpose, in the initial stages of GSoC, it was decided to use the Laplace approximation method implemented in LaplaceRedux.jl as the default Bayesian learning algorithm for ConformalPrediction.jl. This interaction was to be achieved through the development of an interface to the MLJ library.

LaplaceRedux.jl

The first phase of the project was thus devoted to improving the LaplaceRedux.jl package by extending its capabilities. Main result was the creation of an interface with MLJ through MLJFlux. Such interface is currently under review from the MLJ maintainers. Once it is accepted, potential users will be able to use LaplaceRedux.jl through MLJ by loading either the LaplaceRegression or the LaplaceClassification model.

Other significant contributions were the introduction of an option to return distribution objects from Distribution.jl by passing predict_proba = true in the LaplaceRedux.predict function, and the introduction of a new method to evaluate the quality of the uncertainty estimate based on the work of Kuleshov (2018). In particular, two new functions (empirical_frequency_regression and empirical_frequency_binary_classification ) have been introduced for both classification and regression tasks to compute the average empirical frequency per credible interval. A new function called calibration_plot defined in the package TaijaPlotting.jl uses the output of such functions to produce a plot of the empirical frequencies against the predicted probabilities to provide to the users with a quantitative evaluation of how well the neural network is calibrated.

Two additional functions, named sharpness_classification and sharpness_regression, also have been introduced to compute the sharpness of the predicted distributions. An alternative method for recalibration, different from the one already included in LaplaceRedux.jl and based on the work of Laves was then included. Contrary to Empirical Bayes, which tries to estimate the parameters of the prior distribution directly from the data to improve the subsequent posterior estimates, sigma scaling uses a calibration dataset to find a scaling factor $s$ for the variances that maximizes the conditional likelihood without retraining the neural network. In the cases where such conditional likelihood is Gaussian, like in our case, $s$ can be computed analytically and assumes the simple form

$$s = \pm \sqrt{\frac{1}{m} \sum_{i}^m (\hat{\sigma}_{\theta}^{(i)})^2 \lVert y^{(i)} -\hat{y}^{(i)} \rVert }$$

where $m$ is the number of elements in the calibration dataset, $y^{(i)}$ the true expected value and $\hat{y}_{\theta}^{(i)}$ is the predicted result.

Another minor improvement was the introduction of a function that check if the Flux model provided by the user already has a final activation function applied.
For all the new functions introduced in LaplaceRedux.jl, unit tests were created and implemented. The official documentation was then updated to reflect the changes and to show how calibration plots can be used to evaluate the quality of the uncertainty estimates.

Example of workflow

Suppose nn is a neural network trained for a regression task. The Laplace wrapper can be fitted with

la = Laplace(nn; likelihood=:regression, subset_of_weights=subset_w)
fit!(la, data)

The average calibration can be visualized by calling the calibration_plot

using TaijaPlotting
predicted_test_distributions =LaplaceRedux.predict(la, X_test,ret_distr=true)
calibration_plot(la,y_test,vec(predicted_test_distributions),n_bins=20)

which outputs a plot like this:

The sigma scaling factor can then be found by applying the sigma_scaling function over a calibration dataset which was put aside initially, then predicted distributions can be rescaled using the function rescale_stddev

predicted_cal_distributions =LaplaceRedux.predict(la, X_test,ret_distr=true)
sigma = sigma_scaling(vec(predicted_cal_distributions),y_cal)
rescaled_test_distribution= rescale_stddev(vec(predicted_test_distributions),sigma)

Finally, we can reevaluate the level of calibration using calibration_plot once again

calibration_plot(la,y_test,vec(rescaled_test_distribution),n_bins=20)

which shows an improved level of calibration.

The same workflow can be followed for classification tasks. A complete tutorial has been generated and it's available on the official LaplaceRedux.jl documentation webpage.

ConformalPrediction.jl

The second phase of the project was instead devoted to a reorganization of the codebase of ConformalPrediction.jl and to the introduction of a new method called :inductive_bayes_regression to produce conformalized bayes sets as defined in the work of Anastasios N. Angelopoulos, Stephen Bates. Conformalized sets are a type of predictive set constructed to ensure that, under certain assumptions, they achieve a desired level of coverage probability for new data points. They guarantee a marginal coverage probability, meaning that the predictive sets will contain the true value a specified percentage of the time (e.g., 95%) across all instances in the test set, without making any strong distributional assumptions. However, these methods do not adapt directly to varying levels of uncertainty or variability across different regions of the input space; in particular, they can fail to provide the correct coverage when there is heteroscedasticity — that is, when the noise or variance in the data depends on the input features. Conformalized Bayes mitigates this issue by leveraging a Bayesian framework to model the underlying distribution of the residuals and incorporating predictions of the variance as a function of the input. By using the posterior distribution of the model's predictions, conformalized Bayes methods dynamically adjust the width of the prediction intervals in regions where the predicted variance is high or low. This dependence on the predicted variance allows the conformal sets to be narrower in regions with lower uncertainty and wider in regions with higher uncertainty, thereby maintaining more accurate local coverage probabilities while still ensuring valid marginal coverage overall.

More precisely, conformalized Bayes uses as a non-conformity score the negative of the predictive probability distribution estimated by the bayesian regressor; Given a Bayesian Model $\hat{f}(y|x)$ which estimates the value of the posterior distribution of $Y_{test}$ at label $y$ with input $X_{test}=x$ , a conformal score is defined as $$s(x,y)= - \hat{f}(y|x)$$, which is high when the model is uncertain,low otherwise. Once we have the conformal score computed over the calibration data $(y_{cal},x_{cal})$, the quantile $\hat{q}$ corresponding to the level of marginal coverage selected can be computed from the scores. The final prediction sets will then be

$SET(x)= \{y:\hat{f}(y|x) > -\hat{q} \}$

Since LaplaceRedux.jl approximate the posterior predictive distribution as a gaussian this last step can be computed analytically by finding all the values of y so that the condition is satisfied, rather than using a sampling procedure. A function that does exactly this has been implemented in ConformalPrediction.jl.

A similar effort for a conformalized Bayes classifier is still ongoing.

Example of workflow

Given some synthetic data(X,y), once the Laplace model from MLJ (here we assume it is named Lapl) is loaded, it can be wrapped in a conformal model with the option method= :inductive_bayes_regression as conformal method

using ConformalPrediction
conf_model = conformal_model(Lapl, method = :inductive_bayes_regression)
mach = machine(conf_model, X, y)

At this point, the standard MLJ syntax can be used to train the model

MLJ.fit!(mach, rows=train)
data = MLJ.predict(mach,X[test,:])

The following image is an example of a conformalized Bayes set obtained in such a way on for a noisy sinusoidal wave

Project Summary

• Work done

  • Created an interface with the MLJ framework.
  • Added support to Distributions to LaplaceRedux.jl.
  • Made improvements to LaplaceRedux.jl.
  • Updated the documentation.
  • Added new functions to evaluate the quality of the predicted distributions.
  • Added a new method to recalibrate Bayesian neural network based on the work of Kuleshov.
  • Added new unit tests for the new functions.
  • Added new plotting function to TaijaPlotting.jl.
  • Added the conformalized Bayes Regressor to the list of methods supported by ConformalPrediction.jl

• Merged PRs

  • #113 fixed issue with MLJFlux.train
  • #29 updated compat for LaplaceRedux
  • #125 Add support to conformalized Bayes for Regression tasks
  • #103 remaining issue on the mljinterface
  • #99 new glm_predictive_distribution and corresponding changes in LaplaceRedux.predict
  • #25 add plotting functions for the calibration plots and changed TaijaPlotting.jl's plot function
  • #94 Fixed tutorials code snippets

• Work In Progress

  • While waiting for the review of our interface from the maintainers of the MLJ registry, a new branch has been set up to implement whatever changes they will ask us to integrate to make the interface compliant to the MLJ's guidelines.
  • A branch is set up to complete the Conformalized Bayes Classifier.
  • Preliminary work on how to implement the Add-One-In importance sampling among the supported conformal methods in ConformalPrediction.jl is being conducted.

Future work

Due to time constraints and issues with the interface to the MLJ library, I was not able to fully implement the Conformalized Bayes Classifier. However, development is still underway, and a partial implementation is already available, allowing users to utilize a Bayesian neural network to obtain an array of class probabilities. Further work is needed to integrate the conformal prediction framework with the existing Bayesian neural network, optimize the classifier's performance, and ensure compatibility across different MLJ models. Another objective that has been delayed is evaluating whether Turing.jl could be used as a framework for Bayesian models instead of LaplaceRedux.jl. This evaluation will be done in the future to determine which framework offers better flexibility, performance, and integration with the current project setup.

What I learned

From the Google Summer of Code 2024 I have learned tons of useful things:

• Julia Before the GSoC, although I already had experience with Bayesian neural networks with Tensorflow Probability, I had only a limited experience with Julia, but thanks to the GsoC initiative i have learned a lot more over the summer.
• Github Advanced features Before the Gsoc I only used github as a cloud storage for my files and although I already wanted to contribute in some way to some open-source projects I refrained from doing it because I didn't grasp completely all the advanced features of git and github. The GsoC give me the opportunity of learning all these aspects from my mentors. Now finally I have the confidence and the skills necessary to contribute to other projects that interest me.
• Approximate Bayesian inferece with the Laplace Approximation Before the Gsoc, I knew different techniques to create bayesian neural networks, like Concrete Dropout, Deep Ensembles, Hamiltonian Monte Carlo, and Variational Bayes, but I didn't know that the Laplace approximation could be used to obtain the same results at a much lower computational cost. Discovering LaplaceRedux.jl was a nice addition to the list of techniques I understand.

Conclusion and Acknowledgements

I would like to express my gratitude to Google, the Julia community, and my mentors (Patrick Altmeyer and Mojtaba Farmanbar) for this enriching experience. The opportunity to participate to the GSoC 2024 has been both challenging and rewarding and allowed me to enhance my skills and contribute to the julia open-source community. After this fantastic experience, I think I will be part of the open-source community and the Julia community for a long time.

Links:

• Blog: Taija Official Blog