Convergence check for infer functionality

[ThijsvdLaar]

When simulating agents in an online setting, the number of VMP iterations required until (effective) convergence varies over time. An agent with initial vague priors generally requires more VMP iterations than an agent that has already learned an accurate representation. Is there a way to stop VMP iterations once a convergence check has passed?

Ideally, I'm looking for something like

infer(model,
      ...,
      free_energy=true,
      iterations=n_iterations,
      atol=atol)

that stops iterating once $\Delta F = F_k - F_{k-1}$ is (positive and) smaller than atol. Otherwise iterations continue until n_iterations is reached. Is there a way to conveniently obtain this behavior?

[albertpod]

There is no easy mechanism to do so at the moment. You can try out using callbacks for this purpose.
There was a discussion on discourse a while ago about something similar.

[albertpod]

ah, I see both you and @bvdmitri were in that discussion.

[ThijsvdLaar]

A callback would be great indeed. However, the after_iteration callback passes only the FactorGraphModel and the iteration number. If there's a way to pass the free energy values to this callback then that would allow for implementing a custom convergence checker.

[bvdmitri]

Here is what I cam up with @ThijsvdLaar

struct StopEarlyIterationStrategy
    tolerance::Float64
    fe_values::Vector{Float64}
end

function (strategy::StopEarlyIterationStrategy)(model, niteration)
    current_fe_value = 0.0
    # Subscribe on the `BetheFreeEnergy` stream but only `take(1)` value from it
    subscribe!(score(model, Float64, BetheFreeEnergy()) |> take(1), (v) -> current_fe_value = v)
    # Take the previous value from the saved history, use the large value if the first iteration
    previous_fe_value = isempty(strategy.fe_values) ? 10e5 : last(strategy.fe_values)
    # Save the current value in the history
    push!(strategy.fe_values, current_fe_value)
    # Stop early if the previous value is close to the current
    if abs(previous_fe_value - current_fe_value) < strategy.tolerance
        return true
    end
    # Otherwise continue
    return false
end

I've tested this structure on the Bayesian Linear Regression example like this:

strategy = StopEarlyIterationStrategy(1e-3, [])
results = infer(
    model        = linear_regression(length(x_data)), 
    data         = (y = y_data, x = x_data), 
    initmessages = (b = NormalMeanVariance(0.0, 100.0), ), 
    returnvars   = (a = KeepLast(), b = KeepLast()),
    iterations   = 100, # Do 100 iterations
    callbacks = (
        after_iteration = strategy,
    )
)
@show length(strategy.fe_values) # 29 in total is enough
plot(strategy.fe_values)

length(strategy.fe_values) = 29

Note that for this implementation the free_energy = true flag is not necessary, since I query the FE manually.

[ThijsvdLaar]

Thanks @bvdmitri, nice solution for now, works like a charm!

[bvdmitri]

Yeah, we should implement that