Possible Roadmap

This issue outlines possible key features and components planned for the first version of the laplax package. The focus is on providing only the foundational elements necessary for Laplace approximation in JAX.

1. Basic Laplace Function

• Function Signature: laplace(model_fn, params, loss_fn, data) -> (model_fn, params, get_cov_scale)
• Description: This function will be the core entry point for performing the Laplace approximation. It will take a model function (model_fn), parameters (params), a loss function (loss_fn), and a dataset (data). It returns the modified model function, updated parameters, and a function to compute covariance scaling.

2. Covariance Scaling Function (get_cov_scale)

• Function Signature: get_cov_scale(prior, ...) -> cov_scale_mv
• Description: This function will accept parameters such as a prior and return a callable for the covariance-square-root-matrix-vector product (cov_scale_mv). This is essential for calibration and evaluation of the model's uncertainty.

3. GGN Approximations

We aim to support the following Generalized Gauss-Newton (GGN) approximations:

• Full GGN: No approximation.
• Diagonal GGN: Provides a diagonal approximation.
• KFAC (Kronecker-Factored Approximate Curvature): A block-diagonal approximation, where the blocks are given by two Kronecker products.
• Low-Rank: A low-rank approximation method.

Each method will include:

• A specific matrix-vector callable.
• An inverse operation that returns the inverse matrix plus a prior-vector callable.

4. Weight-space uncertainty (cov_scale_mv) to output-space uncertainty (out_cov_scale_mv):

Initial support will include:

• Linearization
• Monte Carlo Sampling

Special consideration is needed for classification tasks, particularly to derive uncertainty for the probit space.

5. Calibration and Evaluation

• Metrics Support: We will support a small set of predefined metric functions.
• Custom Metric Functions: Users can provide custom metric functions that take inputs such as map_prediction, mean_prediction, cov_prediction, and true_target.
• Evaluation Process: These functions will be applied over a specified set of validation or test data points.
• Calibration: Initially, we will support grid search as a method for calibration, requiring the metric to be a single, minimizable function.

Happy to discuss these thoughts below.