Constructing block-banded matrices

[tillahoffmann]

Background: To avoid the XY problem, first a bit of background in case I'm running down the wrong rabbit hole. I am working with state space models of the form $y_t = A y_{t-1} + \epsilon_t$, where $y_t$ is a state vector, $A$ is a transition matrix, and $\epsilon_t$ is Gaussian noise added at each time step. The model is observed through independent observations $x_t=y_t+\gamma_t$, where $\gamma_t$ is Gaussian observation noise.

I would like to evaluate the posterior distribution of $y$ given $x$ which is straightforward (although computationally cumbersome) using the conditioning properties of multivariate normal distributions¹. However, to use the conditioning properties, I need to first evaluate the covariance or precision matrix, i.e., inverse of the covariance matrix.

The precision matrix, has a banded structure because the process is a Markov process. Specifically, the log density is (up to constants independent of $y$)

$$ -\frac{1}{2} \sum_{k=1}^t \left(A y_k-y_{k-1}\right)^\intercal C^{-1}\left(A y_k-y_{k-1}\right), $$

where $C=\mathbb{E}\left[\epsilon\epsilon^\intercal\right]$ is the noise covariance. Only terms involving $y^\intercal_k y_k$ or $y^\intercal_k y_{k-1}$ appear such that the precision matrix looks like

$$ \begin{pmatrix} A^\intercal C^{-1} A + C^{-1} & -A^\intercal C^{-1} & 0 & 0 & \ldots \\ -C^{-1} A & A^\intercal C^{-1} A + C^{-1} & -A^\intercal C^{-1} & 0 & \ldots \\ 0 & -C^{-1} A & A^\intercal C^{-1} A + C^{-1} & -A^\intercal C^{-1} & \ldots \\ &&\ldots&& \end{pmatrix}, $$

where the two terms on the diagonal account for $y_t$ appearing twice: Once in "its own" likelihood and once when being conditioned on by the subsequent state².

The JAX question: Now to the main part of the question: How can I construct this matrix efficiently (including handling batch dimensions). In practice, I have $t\sim300$, $p=2$, and a batch size of $\sim700$. So the output should be a batch of 700 positive definite 600-by-600 matrices.

Attempts: This feels like it should be achievable using a convolution where the rectangular block $\left(-C^{-1} A, A^\intercal C^{-1} A + I, -A^\intercal C^{-1}\right)$ is the filter being convolved with the identity matrix. But, contrary to the standard machine learning setting, now the filters have batch dimensions. I've had a look at jax.lax.conv_general_dilated_patches but couldn't quite get my head around it.

Thank you for the great package and reading this far!

Footnotes

1. We can of course sample the posterior because the density can be evaluated easily (cf. Add RecursiveLinearTransform for linear state space models. pyro-ppl/numpyro#1766), but I'm here interested in the exact distribution. ↩

2. Technically, the last term does not have the $A^\intercalC^{-1} A$ addition, but that small discrepancy can be fixed after constructing the banded matrix. ↩

[tillahoffmann]

Alright, figured it out in case it's useful for anyone. It involves constructing the rectangular block $\left(-C^{-1} A, A^\intercal C^{-1} A + I, -A^\intercal C^{-1}\right)$, padding it with zeros on the right, vmap-ing over shifts for each row, and some reshaping.

def state_space_precision(A, innovation_precision, n):
    """
    Evaluate the state space precision.

    Args:
        A: Transition matrix with shape `(..., p, p)`, where `p` is the dimensionality 
            of the state space.
        innovation_precision: Precision of innovation noise with shape `(..., p, p)`.
        n: Number of steps.

    Returns:
        Precision of the state with shape `(..., n * p, n * p)`.
    """
    # Broadcast arrays and reshape to have a single batch dimension. We'll restore it
    # later.
    A, innovation_precision = jnp.broadcast_arrays(A, innovation_precision)
    *batch_shape, p, _ = A.shape
    A = A.reshape((-1, p, p))
    innovation_precision = innovation_precision.reshape((-1, p, p))
    # Evaluate one of the blocks of the precision matrix which we'll pad and roll.
    offdiag_block = - jax.lax.batch_matmul(A.mT, innovation_precision)
    diag_block = - jax.lax.batch_matmul(offdiag_block, A) + innovation_precision
    row_block = jnp.concatenate([offdiag_block.mT, diag_block, offdiag_block], axis=-1)
    # Pad for rolling, vmap the rolls for efficiency, and discard the part we don't need.
    result = jnp.pad(row_block, ((0, 0), (0, 0), (0, (n - 1) * p)))
    result = jax.vmap(lambda shift: jnp.roll(result, shift * p, axis=-1))(jnp.arange(n))[..., p:-p]
    # Move the rolled dimension to the right position and reshape to get a batch of square matrices.
    result = jnp.moveaxis(result, 0, -3).reshape((-1, n * p, n * p))
    # Set the last element to the innovation precision because there are no subsequent
    # samples.
    result = result.at[..., -p:, -p:].set(innovation_precision)
    # Restore the old batch shape.
    result = result.reshape((*batch_shape, n * p, n * p))
    return result