Skip to content

Support non-reversible quasiseparable kernels#276

Open
burke86 wants to merge 1 commit into
dfm:mainfrom
burke86:fix/nonreversible-quasisep-upstream
Open

Support non-reversible quasiseparable kernels#276
burke86 wants to merge 1 commit into
dfm:mainfrom
burke86:fix/nonreversible-quasisep-upstream

Conversation

@burke86

@burke86 burke86 commented Jul 15, 2026

Copy link
Copy Markdown

Summary

This PR adds support for non-reversible state-space models in tinygp’s quasiseparable kernel machinery.

The previous QSM generator construction works for tinygp’s existing reversible kernels, but it relies implicitly on a reversibility identity that does not hold for general causal systems. In a non-reversible model, the ordering of the stationary covariance and transition matrices matters. Using the reversible ordering can produce incorrectly oriented cross-covariances and disagreement between the quasiseparable, Kalman, and direct dense formulations.

Motivation

I encountered this while implementing a causal filtered-response kernel in which a latent stochastic process drives a downstream response through a one-way linear filter. This structure appears in many settings beyond that particular application, including delayed-response systems, causal multivariate Gaussian processes. These models are generally non-reversible because information propagates from the driver into the response without an equivalent reverse coupling.

Previous behavior

tinygp defines transition_matrix(t1, t2) using an adjoint convention. If $F_{21}$ is the physical forward column-state transition,

$$x(t_2)=F_{21}x(t_1),$$

then transition_matrix(t1, t2) returns

$$A_{21}=F_{21}^{\mathsf{T}}.$$

The previous QSM construction used

$$p_i=h_iP_\infty A_i,\qquad q_i=h_i,\qquad a_i=A_i.$$

For $t_i>t_j$, this represents the lower-triangular covariance as

$$K_{ij}^{\mathrm{old}} = h_iP_\infty A_iA_{i-1}\cdots A_{j+1} h_j^{\mathsf{T}}.$$

Because $A_k=F_k^{\mathsf{T}}$, this is equivalent to

$$K_{ij}^{\mathrm{old}} = h_iP_\infty \left(F_{j+1}\cdots F_i\right)^{\mathsf{T}} h_j^{\mathsf{T}}.$$

The covariance obtained by propagating the column state forward in time is instead

$$K_{ij} = h_iF_iF_{i-1}\cdots F_{j+1} P_\infty h_j^{\mathsf{T}}.$$

The essential difference is therefore

$$\boxed{ P_\infty F^{\mathsf{T}} \quad\text{versus}\quad F P_\infty }$$

For a reversible stationary process, these expressions are equal:

$$F P_\infty=P_\infty F^{\mathsf{T}}.$$

This is why the previous implementation works for tinygp’s existing reversible kernels and why the issue was not exposed by the existing test suite.

For a non-reversible process, such as a one-way driver-to-response system,

$$F P_\infty\ne P_\infty F^{\mathsf{T}}.$$

The previous QSM construction can therefore produce incorrectly oriented cross-covariances. This is especially visible when different observations load different components of a multivariate latent state.

New behavior

This PR stores the physical forward transition in the QSM generators:

$$p_i=h_iA_i^{\mathsf{T}}=h_iF_i,$$ $$a_i=A_i^{\mathsf{T}}=F_i,$$ $$q_i=h_iP_\infty.$$

The resulting lower-triangular covariance element is

$$\begin{aligned} K_{ij} &= p_i a_{i-1}\cdots a_{j+1}q_j^{\mathsf{T}}\\\ &= h_iF_iF_{i-1}\cdots F_{j+1} P_\infty h_j^{\mathsf{T}}, \end{aligned}$$

which is the standard causal state-space covariance.

For reversible kernels, the old and new constructions represent the same covariance. The change therefore preserves existing behavior while extending the QSM machinery to the more general non-reversible case.

Changes

This PR:

  • constructs the symmetric QSM generators using the physical forward-transition orientation;
  • updates scalar kernel evaluation to use the same causal covariance ordering;
  • updates generalized QSM construction used for cross-coordinate matrix multiplication and prediction;
  • documents that transition_matrix returns the adjoint of the physical forward column-state transition;
  • adds a minimal two-state, non-reversible driver-to-response kernel as a regression test.

The test kernel has one-way coupling of the form

$$\text{driver}\longrightarrow\text{response},$$

so its stationary state satisfies

$$F P_\infty\ne P_\infty F^{\mathsf{T}}.$$

It therefore exposes transition-orientation errors that reversible test kernels cannot detect.

Testing

The new tests compare:

  • an explicitly constructed dense state-space covariance;
  • the covariance produced by QuasisepSolver;
  • likelihoods from QuasisepSolver;
  • likelihoods from KalmanSolver;
  • cross-coordinate quasiseparable matrix multiplication;
  • conditioning at the training coordinates;
  • conditioning at new coordinates.

The existing quasiseparable and Kalman tests continue to pass.

Validation on the clean branch:

  • 133 tests passed;
  • 1 optional test was skipped;
  • all pre-commit hooks passed.

User impact

This enables custom quasiseparable kernels for causal and coupled state-space systems where one latent component drives another without reciprocal coupling.

Existing reversible kernels retain their current covariance and likelihood behavior, while non-reversible kernels now agree across the direct dense, Kalman, and quasiseparable implementations.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant