How to specify custom Penalties

This guide shows how the user can specify a customized penalties to help the solver to converge to global minimum.

Motivation

The basic cost in HVI is the negative log of the joint probability, i.e. the likelihood of the observations given the parameters times the prior probability of the parameters.

Sometimes one wants to specify additional knowledge not encoded in the prior, such as one parameter must be larger than another, or entropy-weights of the ML-parameters. For such cases, the solver accept a function that computes additional loss terms. This guide walks through the specification of such additional penalties.

First load necessary packages.

using HybridVariationalInference
using SimpleChains
using ComponentArrays: ComponentArrays as CA
using JLD2
import StableRNGs

This tutorial reuses and modifies the fitted object saved at the end of the Basic workflow without GPU tutorial, that used a log-Likelihood function assuming observation error to be distributed independently normal.

fname = "intermediate/basic_cpu_results.jld2"
print(abspath(fname))
prob = load(fname, "probo");

Write function to compute the penalty loss

The function signature corresponds to the one described in compute_penalty.

In this example we want to avoid local minima when parameter, r1, is larger than 70% of the maximum observation.

# compute the maximum of observed rates at each training site
y_obs = get_hybridproblem_train_dataloader(prob).data[3]
const y_obs_max = map(col -> maximum(x -> isfinite(x) ? x : zero(x), col), eachcol(y_obs))

function compute_penalty_r1(y_pred::AbstractMatrix, addq_pred::AbstractMatrix, 
            θMs_tr::AbstractMatrix{T}, θP::AbstractVector, i_sites_train,
            ϕq::AbstractVector) where T
    penalty = if isempty(i_sites_train) || iszero(i_sites_train[1])
        # for unobserved sites (or test data), assign zero penalty
        n_site = size(θMs_tr,1)
        zeros(T, n_site)
    else 
        # get the maximum of current batch from closure of this function
        y_obs_max_sites = y_obs_max[i_sites_train]
        # add a penalty if r1 is larger than 0.95 times the maximum
        penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- 0.95 .* y_obs_max_sites)
    end
    (; penalty)
end

The penalty is computed for each site in the batch separately. Here, the i_sites_train argument is used to index into precomputed observation maxima.

Update the problem and redo the inversion

HybridProblem provides the penalty_computer keyword argument to specify the Callable that computes the penalty. It defaults to ZeroPenaltyComputer, which returns zero penalty cost.

The argument accepts the function directly (or alternatively construct a CustomPenaltyComputer). Next, the updated problem is solved, using HybridPointSolver and HybridPosteriorSolver.

#prob_pen = HybridProblem(prob; penalty_computer = compute_penalty_r1)
penalty_computer = CustomPenaltyComputer(compute_penalty_r1)
prob_pen = HybridProblem(prob; penalty_computer)

using OptimizationOptimisers
import Zygote
# silence warning of no GPU backend found (because we did not import CUDA here)
ENV["MLDATADEVICES_SILENCE_WARN_NO_GPU"] = 1

# first run a few iterators with updating only optimizing the mean
solver_point = HybridPointSolver(; alg=Adam(0.02))
(; probo) = solve(prob_pen, solver_point; 
    callback = callback_loss(100), # output during fitting
    epochs = 5,
); probo_pen_point = probo;

# starting from this, also estimate the posterior uncertainty parameters
solver = HybridPosteriorSolver(; alg=Adam(0.02), n_MC=3)
(; probo) = solve(probo_pen_point, solver; 
    callback = callback_loss(100), # output during fitting
    epochs = 5,
);

Inspect the computed maxima

The predict_hvi function also evaluates the penalties. Internally, the penalty function is called for each sample, but predict_hvi computes and returns the average for each site.

rng = StableRNGs.StableRNG(112)
n_sample_pred = 200
(; y, θsP, θsMs_tr, ζsP, ζsMs_tr, penalties) = predict_hvi(rng, probo; n_sample_pred);
size(penalties)

The penalties object is a ComponentMatrix, and we can look at a specific site (rows) and a named component (column) returned by

i_site = 3
penalties[i_site, :penalty]

Writing a customized PenaltyComputer

In the above example, the maximum of the observations in the batch are accessed by a global variable, which can lead to type stability and performance problems.

This can be improved. The precomputed maxima can be stored in a struct implementing type AbstractPenaltyComputer with associated function compute_penalty.

struct R1PenaltyComputer{T} <: AbstractPenaltyComputer where T
    r_max::Vector{T}
end
function R1PenaltyComputer(ys::AbstractMatrix)
  r_max = 0.95 .* vec(maximum(ys; dims = 1))
  R1PenaltyComputer(r_max)
end
function HybridVariationalInference.compute_penalty(
    pc::R1PenaltyComputer{T},
    y_pred::AbstractMatrix, addq_pred::AbstractMatrix, θMs_tr::AbstractMatrix, θP::AbstractVector, 
    i_sites_train::AbstractVector, 
    ϕq::AbstractVector
    ) where T
    penalty = if isempty(i_sites_train) || iszero(i_sites_train[1])
        # for unobserved sites (or test data), assign zero penalty
        n_site = size(θMs_tr,1)
        zeros(T, n_site)
    else 
        # get the maximum of current batch from struct
        r_max_sites = pc.r_max[i_sites_train]
        # add a penalty if r1 is larger than the precomputed threshold
        penalty = max.(zero(eltype(θMs_tr)), θMs_tr[:,:r1] .- r_max_sites)
    end
    (;penalty)
end

penalty_computer = R1PenaltyComputer(y_obs)

Rerunning the inversion using with the update PenaltyComputer:

prob_pen = HybridProblem(probo; penalty_computer)
(; probo) = solve(prob_pen, solver; 
    callback = callback_loss(100), # output during fitting
    epochs = 5,
);