Speed up Autoregressive Active Inference example ~40x (49 min -> ~72 s)

[bvdmitri]

Closes #92. Findings I deliberately did not act on are in #93 for @wmkouw.

Advanced Examples/Autoregressive Active Inference took 2956 s (49 min) to build — 6x the next slowest example (Contextual Bandits, 498 s) and the sole reason the examples job timeout was raised to 90 minutes. It now runs end to end, plots and GIF included, in ~72 s.

Behaviour is preserved: all 10 validation seeds still reach 3/3 waypoints.

Where the time actually went

Measured rather than guessed — and it was not where I first assumed. Instrumenting mode(::unBoltzmann) showed 27,810 calls totalling 436 s, essentially the entire runtime. Splitting one call open on an already-compiled run:

corner probe total :     0.04s
optimize() total   :   128.56s
of which f+g evals :     0.73s  (1,558,575 f, 1,558,575 g calls)
Optim overhead     :   127.83s  (99% of optimize)
mean per f/g eval  :    0.236us

Over 99% of the cost was Optim's Fminbox(LBFGS()) scaffolding, not the objective. The configuration outer_iterations=100, iterations=1 rebuilds a barrier problem and a fresh solver state on every outer iteration while making almost no progress per iteration. The 0.2 s time_limit was reached only 3 times in 27,810 calls, so it was never binding either — lowering it would have done nothing.

Two hypotheses I checked and discarded along the way, for the record: it is not a type-explosion problem (a census over a full run finds exactly one concrete energy type) and not primarily compilation (running the loop twice, the second pass is only 1.20x faster).

Changes

Replace the mode optimiser with projected-gradient descent plus a backtracking line search, still seeded from the best box corner so the search stays global. Measured in situ against the old optimiser on identical inputs:

Optim Fminbox total :   130.87s  (14.117 ms/call)
projected gradient  :     0.28s  ( 0.030 ms/call)   -> 463x faster
G(new) - G(old): max=+7.1e-15  mean=-7.989e-01
new strictly better : 5739   equal : 3531   worse : 0

It never returns a worse minimum and finds a strictly better one in 62% of calls, so this is a convergence fix as well as a speed one — see #93 item 5.

Give the expected free energy a closed form. Σ(u) = s(u)/η · V⁻¹ is a scalar multiple of a fixed matrix, so logdet and trace factor into constants; and the regressor x(u) = Pu + q is affine, making s(u) quadratic and μ(u) affine in the 2-D action. Precomputing the coefficients once per message reduces each evaluation to ~20 flops with an analytic gradient, removing ForwardDiff from the hot path.

Verified against the original expression across 300 random problems at the real dimensions:

optimized helper vs original rule : max rel err = 1.110e-15
analytic gradient vs ForwardDiff  : max abs err = 3.553e-15

value    : original 2.62 us/call, optimized 0.056 us/call  -> 47x
gradient : ForwardDiff 55.49 us/call, analytic 0.172 us/call -> 323x

Stop inverting the same matrix twice. The rules did Λ = inv(Λ) and then inv(Λ) again inside the energy — recovering the original matrix on every evaluation. posterior_predictive already documented the intended form (x'*U*x, with U used directly).

Hoist mode(...) out of the energy closures. They do not depend on the action, and several arguments are themselves unBoltzmann, so each was triggering a nested optimisation per evaluation.

Carry the gradient through prod, so products of energies keep an analytic gradient rather than falling back to AD.

Parameterise unBoltzmann — G::F, N::I, D::R instead of ::Function, ::Integer, ::Rectangle (the last was also abstract: Rectangle{T} <: HyperRectangle{T}).

Readability

The 42 @rule MARX definitions had only ~5 distinct bodies — 10 :in rules byte-identical apart from argument names, likewise 8 :outprev1 and 5 :out, plus 18 that were just return Uninformative(). Each group now delegates to one documented helper, with the rules reduced to thin wrappers. All 42 signatures are kept, so dispatch coverage is unchanged.

Dimensions and action limits now reach the rules as MARXMeta rule metadata via @meta instead of untyped globals read from a later cell, and the simulation loop is wrapped in run_simulation rather than running at top level under global.

Note that attaching meta is all-or-nothing: a rule declaring no meta:: argument generates a method with meta::Nothing and stops dispatching once meta is attached (ReactiveMP/src/rule.jl:478) — hence all 42 wrappers take it.

Validation

seed	waypoints	switches at
1	3/3	[103, 130, 145]
2	3/3	[102, 127, 150]
3	3/3	[103, 129, 144]
4	3/3	[103, 129, 145]
5	3/3	[103, 129, 144]
6	3/3	[103, 140, 157]
7	3/3	[103, 129, 144]
8	3/3	[103, 129, 145]
9	3/3	[103, 140, 158]
10	3/3	[103, 129, 144]

Matching the bar set in #80 ("validated on 10 random seeds: all reach all 4 waypoints within 400 steps").

The GIF is regenerated: better-converged planning reaches the waypoints sooner, so the trajectory legitimately differs from the published one. @wmkouw — flagged for your sanity check as #93 item 5.

Not changed

Two sign conventions in the rules look inconsistent but turn out to be a compensating pair (flipping both is bit-identical; flipping either alone gives 0/3 waypoints), and the mutualinfo helper disagrees with the rules in a way that affects only the EFE landscape figure. Both are written up in #93 rather than changed here, so this PR stays purely about performance.

🤖 Generated with Claude Code