Codex/ivashchenko filter - #314
Open
thorek1 wants to merge 16 commits into
Open
Conversation
A pruned second-order solution is linear in the augmented state
z = [x1 ; x2 ; x1[past] ⊗ x1[past]]
(the pruned state-space representation of Andreasen, Fernández-Villaverde &
Rubio-Ramírez), and the observation is a plain selection from it, so a Kalman
filter applies. That is the quadratic Kalman filter of Monfort, Renne &
Roussellet (2015).
The transition is exactly linear in z and the conditional first and second
moments are closed form. Writing aug1 = ā + Sε, every block of the innovation
is w = Gε + H(ε⊗ε − vec(I)); the Gaussian third moment vanishes, so
Var(w) = GG' + H(I+K)H' with K the commutation matrix. H is constant, G is
state dependent and evaluated at the filtered mean each period. What is
approximated is the conditional *distribution* — the innovation is quadratic in
ε, so the recursion is the best linear projection rather than the exact
conditional mean — and the Kronecker block is treated as a free state rather
than constrained to equal the square of the first.
Verified in three steps of increasing strength, because the obvious test is the
weakest: on a linear model 𝐒₂ = 0 leaves the quadratic blocks inert, so
agreeing with the Kalman filter there (it does, exactly) proves only the
plumbing. So the augmented transition is also checked against a Monte-Carlo
evaluation of the package's own pruned recursion, and the filter against a
particle filter on a genuinely nonlinear model, where they agree to 0.05 log
points at a measurement-error variance of 1e-4 and both approach the inversion
filter's zero-measurement-error limit.
Initialisation solves (I − 𝒜)z̄ = c directly and the ergodic covariance by
doubling; iterating either would need thousands of steps on a model with roots
near unity, which cost 0.23 log points before it was fixed.
Co-Authored-By: Claude <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
Adds `filter = :quadratic_kalman` to the registry and routes it in `get_loglikelihood`. It is gated to `algorithm = :pruned_second_order` — the only case in which the augmented state space is linear — and falls back to the inversion filter with a message elsewhere rather than erroring. Smoothing is turned off, since no smoother is implemented. Missing observations are rejected explicitly. The implementation is now type generic, so forward-mode AD flows through the closed-form moment algebra; there is no finite differencing anywhere inside the filter. The gradient is checked against central differences on the likelihood (agreement to 1.3e-7). A hand-written reverse-mode rrule is *not* implemented — at 2·nVars + nPast² states that is a substantial separate piece, and forward mode costs one solve per parameter. SW07 accuracy, measured against a particle filter at matching per-observable measurement error on the EA data and initial values from the nonlinearities repository (the parameters at which the inversion filter is well behaved at second order): T=30 QKF -382.87 PF -406.10 (sd 1.02) gap 0.77/period T=60 QKF -748.15 PF -787.98 (sd 2.91) gap 0.66/period T=138 QKF -2150.55 PF -2231.13 (sd 2.73) gap 0.58/period so the filter overstates the likelihood by roughly 0.6 per period there, against 0.001 per period on the small RBC. The difference tracks the strength of the nonlinearity: max|𝐒₂| is 1451 on SW07 against 2.5 on the RBC. That is the linear-projection approximation showing, not an implementation error — the augmented transition is verified against a Monte-Carlo evaluation of the package's own pruned recursion. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
…e it ~7x faster
Reverse mode. The recursion is split out as `quadratic_kalman_recursion` and given
a hand-written `rrule`. The derivation is made tractable by noticing that G(z) is
*affine* in z — both the 𝐒₂ block and the Kronecker block are linear in ā, which
is affine in z — so vec(G) = g₀ + Λ(Pz) and the whole state dependence of the
innovation covariance collapses to one matrix. That removes the per-period
Kronecker adjoints entirely; Λ is built exactly, column by column, from the affine
map rather than by differencing.
All nine cotangents (𝒜, c, QH, g₀, Λ, Hm, data, z₀, Σ₀) agree with ForwardDiff to
~1e-16, and that comparison is now a test. One subtlety cost a wrong derivative
before it was found: the forward pass symmetrises F, so the cotangent reaching CP
and Hm is (F̄+F̄')/2 — omitting that left d/dHm wrong by 2% while every *other*
derivative still looked exact.
The rule covers the O(T·nz³) recursion, which dominates and scales with the
sample. Building the system matrices is O(1) in T and is left to ordinary AD, so
the two compose normally rather than nesting AD inside a pullback.
Speed. Two exact reformulations, neither changing the likelihood (-2150.547
before and after on Smets-Wouters):
* q = x₁ₚ⊗x₁ₚ = vec(x₁ₚx₁ₚ') is symmetric, so carry vech instead of vec via
duplication/elimination matrices: the Kronecker block drops from nPast² = 729
to nPast(nPast+1)/2 = 378.
* only the past states and the observables are ever read out of the x₁/x₂
blocks, so retain those rows instead of all nVars: 67 rows become 34.
Together nz goes from 863 to 446 and the filter from 7.09 s to 1.02 s per
evaluation — 7x, against 7.24x predicted by the cubic scaling.
Not done: the top-level `get_loglikelihood` rrule does not dispatch to this path,
and a pre-existing guard refuses reverse-mode AD whenever measurement error is
active, so `Zygote.gradient` cannot yet reach it end to end. ForwardDiff works.
Co-Authored-By: Claude <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
The recursion already had a verified hand-written adjoint, but nothing could reach it: the top-level rrule looks for `rrule(calculate_loglikelihood, Val(filter), …)` and silently returns a zero gradient when none exists. Three things were missing. 1. The build adjoint. Cotangents of (𝒜, c, QH, g₀, Λ) are now pushed back onto 𝐒₁/𝐒₂ analytically, including the ergodic initialisation: z₀ = (I−𝒜)⁻¹c gives a transposed solve, and Σ₀ = 𝒜Σ₀𝒜' + Q₀ gives a second Lyapunov equation X = 𝒜'X𝒜 + Σ̄₀, after which Q̄₀ = X and 𝒜̄ += 2X𝒜Σ₀. Verified against ForwardDiff at ~4e-15 for ∂𝐒₁, ∂𝐒₂ and ∂data. 2. The filter now routes through `calculate_loglikelihood`, with an rrule at that interface returning ∂𝐒 in the expected position, scattered from the retained rows back onto the full solution matrices. 3. The reverse-mode guard refused any likelihood with measurement error. It is relaxed for this filter only, since its adjoint covers H. One wiring bug worth recording: the top-level rrule never passed `measurement_error` to the inner rrule. That was harmless while the guard made measurement error impossible, but here it meant reverse mode differentiated the H = 0 likelihood while the primal used H > 0 — a *finite* gradient that was simply wrong (relative deviation 2.5). Zygote and ForwardDiff now agree to 6e-14, and that comparison is a test. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
Measured before guessing: the filter spent 60% of its time not in the covariance
propagation but in allocation — 2509 MiB over a 138-period Smets-Wouters run,
roughly 18 MB per period, mostly nz×nz temporaries from `𝒜*Pc*𝒜'`, `G*G'` and
the `𝒞` products.
Three changes:
* every buffer is allocated once per call and reused, with `mul!` throughout;
the rank-nExo term G G' folds into Pp as a single gemm update rather than
forming Q separately;
* 𝒞 is a selection matrix, so `𝒞*zp`, `𝒞*Pp` and `CP*𝒞'` become indexing
instead of three gemms;
* the symmetrisations write both triangles in one pass rather than building
(X+X')/2 as a fresh matrix.
0.986 s → 0.485 s and 2509 MiB → 171 MiB on SW07. The covariance propagation is
now 92% of the remaining runtime, so further gains have to come from that term.
One trap this surfaced: preallocation fixes the element type, so the promotion
has to cover *every* differentiable argument. Covering only 𝒜, Y and g₀ left
forward-mode AD working with respect to those three and failing with respect to
c, QH, Λ, Hm, z₀ and Σ₀ — caught by the cotangent tests.
Co-Authored-By: Claude <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
Adds a full section to the Filters page: the augmented-state construction and why pruning is what makes it possible, what is exact versus approximated, the source and size of the likelihood bias, when the filter is and is not usable, and a measured cost breakdown. The two findings worth recording: The bias is not what it looks like. It does not come from 𝐒₂ — zeroing its ε⊗ε block leaves rank(Q) unchanged. It comes from kron(V,V) in the *first-order* solution: x₁[past] = (deterministic) + Vε, so q = x₁[past]⊗x₁[past] inherits Vε⊗Vε whatever 𝐒₂ is. The quadratic noise is intrinsic to carrying a Kronecker term as a state. And because the observation has zero loading on q, the data can never shrink the resulting fictitious uncertainty, so the error persists as the measurement error goes to zero (converging to -1.52 rather than to 0). But it is almost a level shift, so it is far less damaging than it first looks: profiled against the exact inversion likelihood the gap varies by only ~0.2 log points across a parameter grid and the mode is unchanged. Latent states — the filter's actual purpose, per Kollmann (2015) — are good: 2.8% relative RMSE on capital, ~11% on the unobserved shock processes. So it is fine for state estimation and point estimation, and unsafe for model comparison, where the level error differs across models with the shock-to-observable ratio. Cost breakdown measured per period at nz=446: the two nz³ triple products are 91% of the loop (1.09 ms and 1.26 ms), everything else 9%. Against the inversion filter's n_ε×n_ε solve that is 446³ vs 7³ — about 2.6e5 in flops. Structural, not tunable; sparsity measures 10× slower since 𝒜 is ~50% dense. Also records that this is Kollmann's filter, not the Monfort-Renne-Roussellet quadratic Kalman filter: theirs has a linear transition and a *quadratic measurement*, so the data loads on the Kronecker block directly and shrinks the very uncertainty that is irreducible here. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
`run_quadratic_kalman` carried its own inlined doubling loop with looser tolerances (1e-12, 60 iterations) than the `qkf_lyapunov` helper used by the pullback (1e-15, 80), so the forward and reverse passes initialised at slightly different ergodic covariances. Both now call one helper, which dispatches Float64 problems to `solve_lyapunov_equation` — the same workspace-backed doubling solver the pruned second-order moments code uses — and keeps the self-contained loop as the fallback for AD element types, which that solver does not accept. The forward pass hands its converged covariance to the reverse pass, which was re-solving the identical equation; passing it as `initial_guess` turns that second solve into a residual check (46.6 ms -> 3.1 ms on SW07). Measured on SW07 at pruned second order, 7 observables: 485 ms -> 425 ms per likelihood, with the log-likelihood unchanged to the last digit (-2264.4493051268087) and all 31 quadratic-Kalman tests passing, including the hand-written rrule cotangents against ForwardDiff. Warm-starting across parameter draws was measured and deliberately not done: the guess only pays when it is exact, and a 1e-6 relative parameter move already makes it a net loss. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
The quadratic Kalman section explained the filter on its own terms but never put it next to the linear one, which is where the intuition actually is: it is the same recursion, and setting S2 = 0 collapses one onto the other exactly. Adds a side-by-side of the two loops as they appear in the source, a table of the structural differences, and short notes on the three that carry consequences — the noise covariance becoming state-dependent (the conditional heteroskedasticity a second-order solution adds), the innovation ceasing to be Gaussian (why the linear filter is exact and this one is not), and the cost being cubic in a squared dimension (n_past^6), which governs when the filter is usable. Dimensions are the measured Smets-Wouters ones: 34 retained rows against an augmented 446 = 2*34 + 27*28/2. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
The third-order analogue of the quadratic filter, built on the same fact: pruning truncates the Kronecker hierarchy at a fixed rung at every order, so the pruned third-order solution is exactly linear in the augmented state [x1; x2; x3; a*a; a*b; a*a*a]. The substance is the closure. Writing a_n = M a + v with v state-independent, each new Kronecker block resolves back onto blocks the state already carries — q11' onto q11, q12' onto q12 and q111, q111' onto q111 and q11 — and no fourth-order block appears because a_n has no q11 component. Computing the new blocks the obvious way, as kron(a_n, a_n), is quadratic in z and destroys the linearity the filter rests on without failing loudly, so the affineness of the step is asserted directly in the tests. Validated on an RBC model (2 shocks, 3 past states): the step is affine to 3e-15, reproduces the pruned third-order recursion to 2e-19 on every block including the Kronecker ones, its quadrature moments match Monte Carlo, and the log-likelihood matches a converged bootstrap particle filter at 181.78 against 181.43 over 60 periods — 0.006 per period, better than the quadratic filter's 0.025 on the comparable model. It only fits small models. The augmented dimension is 3nr + 2nPast^2 + nPast^3 and the covariance recursion is O(nz^3), so cost grows as nPast^9: fine to nPast=10, marginal at 12, hopeless at Smets-Wouters' 27 (nz = 21243, 3.6 GB per matrix). The build refuses above CUBIC_KALMAN_MAX_DIMENSION rather than appearing to hang. Not yet done, and noted in the docs: analytical assembly instead of quadrature, vech compression of the Kronecker blocks, and a hand-written rrule. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
Four of them, all measured on the RBC test model (2 shocks, 3 past states) with the log-likelihood unchanged to the last digit at every step (181.77600836054...): - vech compression of the symmetric Kronecker blocks. q11 and q111 are symmetric and fully symmetric respectively, so the state carries one entry per sorted multi-index. Applied by index maps rather than duplication/elimination matrices, which here would cost more than the compression saves. nz drops 60 -> 40 on the test model and 21243 -> 4863 at Smets-Wouters size; since the recursion is O(nz^3) this is two orders of magnitude on a mid-sized model, and it moves the practical wall from nPast ~= 12 to ~= 20. - An allocation-free step. The step does a few hundred flops but was allocating ~13 kB per call, and it runs (nz+1)*n_nodes times to build the transition plus n_nodes times per period — so it was entirely allocation-bound. With preallocated buffers it is 4.33 us/12960 B -> 2.71 us/0 B. The symmetric output blocks are written straight to their canonical slots, so the full nPast^3 vector is never formed. - An allocation-free recursion on preallocated buffers with in-place BLAS, and the observation applied by indexing its three selected rows instead of a gemm with a 0/1 matrix. - The quadrature contracts its node evaluations with one gemm rather than n_nodes rank-one updates, and the transition build skips the variance it never uses. End to end: 9.1 ms -> 5.9 ms on the test model, where the covariance recursion is only ~13% of the time and the quadrature dominates. The gains are far larger on models where nz is big enough for the O(nz^3) term to matter. Tests grow to 16: the compression maps are checked to round-trip a genuine symmetric Kronecker product, and the transition reference now forms every Kronecker product in full before compressing, so it exercises the compressed algebra rather than assuming it. Still not carried over: analytical assembly of the transition in place of quadrature, and a hand-written rrule. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
The in-place step writes the row-major flattens out by indexing, so nothing calls it any more; only the comments still use the term, which is now defined where they can see it. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
f(z, ·) is a polynomial of degree <= 3 in eps whose coefficients are affine in z. Recovering that coefficient matrix C(z) once gives both moments in closed form: E[f] = C(z) m, m_a = E[eps^a] Var(f) = C(z) Psi C(z)', Psi_ab = E[eps^(a+b)] - E[eps^a] E[eps^b] and since the shocks are independent standard normals, E[eps^a] factorises into double factorials, so Psi is a closed form rather than a set of Isserlis pairings. C(z) is recovered by interpolation on C(nExo+3, 3) points, which is also where the tensor Gauss-Hermite rule is left behind: its node count grew as npt^nExo (16384 for seven shocks at npt=4) against 120 for the coefficient basis. Consequences: no quadrature per period at all — Q(z) is one matvec and two gemms — and the build no longer blows up in the number of shocks. The quadrature path is kept and the analytic assembly is tested against it rather than assumed: A and c agree to 1e-9, Q(z) to 1e-8 relative at a non-trivial z, and the monomial moment vector matches tensor Gauss-Hermite to 1e-10. The log-likelihood is unchanged at 181.776008360546. 5.9 ms -> 2.3 ms on the test model (9.1 ms before this round of work started). Tests 16 -> 21. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
Forward mode now works end to end: the element type is threaded through the system build, the step workspace and the recursion, so ForwardDiff duals flow through. Gradients match central finite differences to 7e-11. The in-place BLAS path is kept for Float64 and the AD path falls back to allocating solves, since rdiv! has no Cholesky method for dual element types. That is the right mode for this filter: it is capped at nPast ~ 20 by its own cost, so a gradient is a handful of few-millisecond primal passes. More important is what reverse mode was doing. There is no hand-written rrule for this filter, and the top-level rrule's `llh_rrule === nothing` branch responds to that by returning on_failure_loglikelihood with an all-zero gradient — no error, no warning. Measured on the test model: Zygote returned exactly zeros where the true gradient has entries up to 3e4. A sampler would have run on that and produced garbage silently. Reverse mode now errors and names forward mode as the alternative. The existing measurement-error guard hid this in the common case, so the test pins the new guard specifically by asking for a gradient *without* measurement error, where that older guard cannot be what fires. Tests 21 -> 24. Quadratic Kalman tests still 31/31 after the rrules.jl change. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
Reverse mode now works and matches central differences to ~1e-10, on every
measurement-error shape (none, diagonal, full covariance). The chain
S -> sys -> {f(z_k, eps_p)} -> (A, c, c0, Lambda) -> llh is differentiated in
three pieces, each verified against ForwardDiff *in isolation* so a regression
localises instead of just moving the end-to-end number:
- the step adjoint (5e-16). Everything the step builds from z and eps alone --
aug, K2, K12, K3, the Q blocks -- is constant here, so only the paths through
the solution matrices carry cotangents.
- the build adjoint (2e-15). The build is linear in the collected step
evaluations, so those maps transpose directly and the work is replaying the
step adjoint over the same (nz+1)*N points the forward pass visited.
- the recursion adjoint, mirroring the quadratic filter's verified one; the
structural difference is Q = C Psi C' in place of GG' + Q_H, whose cotangent is
2 Q_bar C Psi because P_bar_p is symmetrised before use.
Reverse mode is now the better choice here as well as the correct one: 12.0 ms
against a 2.3 ms primal and independent of parameter count, where forward mode is
44.3 ms for seven parameters and grows linearly.
Also fixes a 6.5x primal regression introduced with forward-mode support. The
promoted element type was stored as a field of the system, so `Matrix{sys.Tv}`
inferred as DataType rather than Type{Float64} and lost specialisation
throughout; the hot paths now read `eltype(sys.S1)`, which is inferrable, and the
primal is back to 2.27 ms from 14.94 ms.
The zero-gradient guard added in the previous commit is removed, since the
fallback it protected against is no longer reachable.
Tests 24 -> 28. Quadratic Kalman still 31/31 after the rrules.jl changes.
Co-Authored-By: Claude <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
Profiling the build showed it is (nz+1)*N step evaluations and essentially nothing else, and that inside a step one operation dominated: contracting the Kronecker input K3 against S3. That matrix is very wide and very sparse -- 8x1331 on a four-shock model, 23% dense -- so the product is memory-bound, and it measured 5.5 of the 10.8 us a step took. Only the structurally nonzero columns of S2 and S3 are now kept (536 of 1331 for S3, 65 of 121 and 40 of 66 for the two S2 paths), which shrinks both the vector that has to be built and the product that consumes it, and drops those iterations from the branchy K3 loop as well. Liveness is taken from the stored pattern of the *sparse* solution matrices, not from numerical zeros of the densified copy. A column that merely happens to vanish at one parameter draw may be nonzero at the next, and dropping it would silently zero a real derivative rather than fail. The existing exactness test guards this directly: it compares the restricted step against a reference built from full Kronecker products. Measured on a four-shock, six-past-state model (nz = 137): step 10.8 -> 6.4 us build 56.2 -> 34.0 ms step pullback 26.3 -> 18.0 us The build's pullback also collapses two sums that were being accumulated one column at a time, each allocating an nz-by-N array per pass. Checked what else sparsity could buy and it does not: Psi is 9% dense and block-diagonal by monomial parity, but it appears only in the smaller of the two products forming Q, and its rank is N-1, so factoring removes just the constant monomial. Also validates the reverse-mode adjoint on a second, wider model: Zygote against central differences is 1.1e-9 there, with the gradient at 4.8x the primal. Co-Authored-By: Claude <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019YdYHujHfAMxP3jt9kwzzT
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## particle-filter #314 +/- ##
===================================================
- Coverage 84.16% 79.68% -4.49%
===================================================
Files 36 39 +3
Lines 31696 33493 +1797
===================================================
+ Hits 26678 26690 +12
- Misses 5018 6803 +1785 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.