A generalization of the Feynman-Kac representation of state-space models

[AdrienCorenflos]

Feynman--Kac models are usually represented as a pair $M(x_t \mid x_{t-1})$, $G(x_{t-1:t})$ which is somewhat very tied to the SMC implementation and not super compatible with other (Gaussian-approximated for example) perspectives.

It does make sense to generalize the structure, and I think the perspective I took in my PhD thesis (Section 5.2.3), which represents them in the same way as what Simo and Angel did in their paper:

Of course, we can't really compute these things in general, but we can say how they should be computed.
For instance, the kernel $F(x_t \mid x_{t-1})$ and function $H(x_{t-1})$ can be simulated and computed together.
I'm not sure what the best representation would actually be, but the idea is that we can really just represent objects by the way they should be computed rather than by their probabilistic structure.

[SamDuffield]

Just adding the extract from Särkka and Garcia-Fernandez.

So the idea would be that the methods would have the following structure

• Filtering state, NamedTuple encoding $a_k = (f_k, g_k)$ sufficient to define $(p(x_k \mid y_k, x_{k-1}), p(y_k \mid x_{k-1}))$
• Filtering binary associative operator, function encoding $a_i \otimes a_j = (f_i, g_i) \otimes (f_j, g_j)$
• Smoothing state $a_k$ sufficient to define $p(x_k \mid y_{1:k}, x_{k+1})$
• Smoothing binary associative operator $a_i \otimes a_j $

[SamDuffield]

Ok need to think further about this but could be really nice.

I think discrete and generalised Kalman should fit, with SMC the main thing to think about.

I would be inclined to not abstract away the observation $y_k$ since linear Gaussian methods really require it. Maybe a log-likelihood of the form $G(x_{k-1}, x_k, y_k)$ is the way to go.

[AdrienCorenflos]

Yeah but then the operation does not need to be associative (won't be in SMC). It's really just a question of representation

[SamDuffield]

Been thinking about this a bit and wondering if something like the below makes sense. Although I'm not quite sure how to handle inputs and observation in a way that is general and works with associative_scan where appropriate.

from typing import NamedTuple, Protocol
from cuthbert.types import (
    ArrayTree,
    ArrayTreeLike,
    KeyArray,
)

# State encoding parameters needed to compute
# f(x_k | x_{k-1}) = p(x_k | y_k, x_{k-1}, inputs)
# and
# g(x_{k-1}) = p(y_k | x_{k-1}, inputs)
FilteringState = ArrayTree

# State encoding parameters needed to compute
# k(x_k | x_{k+1}) = p(x_k | y_{1:k}, x_{k+1}, inputs)
SmoothingState = ArrayTree


# Initiate first filtering state encoding parameters needed to compute
# f(x_0 | None) = p(x_0 | inputs)
# and
# g(None) = 1
class Init(Protocol):
    def __call__(
        self,
        inputs: ArrayTreeLike,
        key: KeyArray | None = None,
    ) -> FilteringState: ...


# Binary operator (could be associative but not necessarily)
# Combines filtering states in a way that is consistent with the filtering distribution
class FilteringOperator(Protocol):
    def __call__(
        self,
        state_prev: FilteringState,
        state: FilteringState,
        observation: ArrayTreeLike,
        inputs: ArrayTreeLike,
        key: KeyArray | None = None,
    ) -> FilteringState: ...


# Binary operator (could be associative but not necessarily)
# Combines smoothing states in a way that is consistent with the smoothing distribution
class SmoothingOperator(Protocol):
    def __call__(
        self,
        state_prev: SmoothingState,
        state: SmoothingState,
        inputs_prev: ArrayTreeLike,
        inputs: ArrayTreeLike,
        key: KeyArray | None = None,
    ) -> SmoothingState: ...


class SSMInference(NamedTuple):
    init: Init
    filtering_operator: FilteringOperator
    smoothing_operator: SmoothingOperator

[AdrienCorenflos]

I was thinking about this rough structure:

import jax
import jax.numpy as jnp
from jax.scipy.stats import norm

from typing import NamedTuple, Any

class BootstrapState(NamedTuple):
    """State of the PF at time 
    """

    key: Any
    weights: Any
    particles: Any
    genealogy: Any

    inps: Any

    def f_rvs(self, key, xs):
        return xs + jax.random.normal(key, xs.shape)

    def g_rvs(self, xs)
        return norm.logpdf(self.inps, xs)

    @classmethod
    def init_as(cls, key, dist):
        ...



class BootstrapOperator(NamedTuple):
    """Operator for the PF
    """

    # Some hyperparameters
    resampling: Any = lambda k, w: jax.random.choice(k, w.shape[0], w.shape, p = w / w.sum())

    def combine(self, state_1, state_2):
        key_resampling, key_propose = jax.random.split(state_1.key)
        ancestors = self.resampling(key_resampling, state_1.weights)
        next_particles = state_1.f_rvs(key_propose, state_1.particles[ancestors])
        next_weights = state_2.g_rvs(next_particles)
        return BootstrapState(state_2.key, next_weights, next_particles, state_1.genealogy)

[Sahel13]

I guess what I don't like is the user having to prepopulate weights, particles, and genealogy, which are the algorithmic outputs.

[AdrienCorenflos]

Tbh, I think it's a byproduct of working with JAX more than anything. True online I a bit out of scope.

[AdrienCorenflos]

Also, they probably wouldn't need to if we gave some constructor for it that just populates it as NaNs

[AdrienCorenflos]

Yeah I understand that for smoothing, we do need this. But the filtering algorithm could output the states at every time step without the user having to pre-specify them. For particle filters, to me it feels cleaner and more intuitive to do it the usual way, having the initial distribution, transition function, and log potential function as the inputs, and have the algo do the rest.

As for our KF implementation, we have to provide all the matrices for every time step at the beginning, but only because it's necessary for the parallel algo.

Technically this is clunky too then, and mostly driven by the parallel line.

[SamDuffield]

I'd be happy with init and combine functions that behave something like the following (unified across inference methods)

Online filtering:

filter_state = init(inputs[0])

for t in range(1, T):
    init_state_t = init(inputs[t])
    filter_state = combine(filter_state, init_state_t)

Offline filtering:

init_states = vmap(init)(inputs)

def body_fun(prev_state, init_state_t):
    filter_state = combine(prev_state, init_state_t)
    return filter_state, filter_state

filter_states = scan(body_fun, init_states[0], init_states[1:])[1]   # If appropriate can use associative_scan

Noting that we might want to change the names of init and combine as well as have an explicit observation argument (useful for Kalman).

Although I'm not sure this pairs too well with the generalized_kalman idea as that requires a state to linearize around (which you won't have in init). Maybe thought there's a way to do it that still supports associative scan in the case the model is linear-Gaussian.

[SamDuffield]

Unified protocol

• We can have combine and associative_combine (the latter can raise NotImplementedError)
• For linearized Kalman, the state can own the linearization
• Then build function to generate unified Inference NamedTuple

[SamDuffield]

Some more ramblings from our meeting

def build(resampling):
     return Inference(combine=partial(combine, resampling=resampling))


def combine(state_1, state_2, resampling):
    key_resampling, key_propose = jax.random.split(state_1.key)
    ancestors = resampling(key_resampling, state_1.weights)
    next_particles = state_1.f_rvs(key_propose, state_1.particles[ancestors])
    next_weights = state_2.g_rvs(next_particles)
    return BootstrapState(state_2.key, next_weights, next_particles, state_1.genealogy)
    
 
inference = pf.build(resampling)
inference = extended_kalman.build(covariance_functions)

state = inference.init(inputs[0])
state = inference.filter_combine(state, inputs[1])

[SamDuffield]

Closing this discussion as we have aligned on an interface