UnHiPPO: Uncertainty-aware Initialization for State Space Models

Thank you for your review and careful reading.

Phrasing

We will clarify Section 4 based on your feedback. The phrase "observations do not exist in HiPPO" refers to the fact that the data does not take the role of an observation in HiPPO, but rather of a control signal. We will update the section to make this clear. We have also adopted your clarification "best possible compression of $f$ in the $L^2$ space".

Effect of noise level

https://figshare.com/s/132436b00e91612513fd

We ran two experiments similar to Figure 7, but this time the noise level $\rho^2$ is set to a lower and a higher value. The best results were obtained for $\sigma^2 = 10^{10}$. However, as we increase the noise level $\rho^2$, we notice that higher $\sigma^2$ values perform comparatively well, as opposed to what happens in Figure 7, when $\sigma^2 = 10^{14}$.

What is the effect of ignoring the time-variance of the true dynamics?

We examined the effect of this in Figure 5 and the surrounding text in the paragraph "Time-invariant Dynamics" empirically, but we did not analyze it theoretically.

How can we select $\sigma^2$?

At the moment, our best strategy for choosing $\sigma^2$ is empirical. The reason that $\sigma^2$ cannot be chosen as the noise in the data comes down to $\Sigma = I$ in the Kalman filter. Because this adds uncertainty to all degrees of the polynomial representation at each step, it also increases the uncertainty in the high-degree components. Even thought the dynamics are regularized, the highest degrees still grow quickly, so $\sigma^2$ needs to be of a similar order of magnitude as $B_H^T P_k^{-}B_H$ in $s_k$ to have an effect. Choosing to the diagonal of $\Sigma$ to fall quickly should be able to counteract this effect, though it is unclear how quickly it would need to fall exactly.

Can UnHiPPO be thought of as a regularization scheme, where $\sigma^2$ controls the magnitude of regularization?

Yes, as we have shown in Figure 4, $\sigma^2$ controls directly how sensitive the system is to high-frequency components / noise in the data.

Can UnHiPPO be extended to HiPPO variants other than LegS?

We decided to focus on the LegS variant of HiPPO, because it is the most relevant one in the literature. Extending the approach to other polynomial bases should be possible, because we only use properties of Legendre polynomials explicitly in Section 4.1 to derive the regularized HiPPO matrix. The same derivation should be theoretically possible in other bases, thought it might turn out to be either unnecessary if the basis behaves well under extrapolation or not produce a closed-form of $Q_i$.

Thank you for your review. Please see Figure 9 in Appendix B for a visualization of the effect of different discretizations. It demonstrates the instability of some methods in both HiPPO and UnHiPPO and in particular the remarkable stability of the closed-form solution.

Thank you for your review.

How was the noise added? How were the signals normalized?

We follow LSSL implementation and the code from repository of the authors. After the audio files are loaded, signals are divided by 32k to be normalized (see here). After the noise is added, we apply z-score normalization.

What type of noise was added?

The two noise sources in equations (7) and (8), namely $\mathrm{d}\mathbf{\beta}$ and $\varepsilon_{t_k}$, are Gaussian. In line with the theory, we sampled Gaussian noise with 0 mean and fixed variance. We have not experimented with other noise sources.

Computation of parameters:

P. 6, paragraph following (29): “In contrast to LSSL where we can get the discretized dynamics at any t directly, for UnLSSL we compute them for all integer steps $t ∈ [t_{max}]$ and then select a subset.” – I assume that this is necessary since the uncertainty-aware transition matrix and input vector are obtained from the discrete-time update equation in (24). Please can you clarify this point?

That is correct. In the end, the parameters are tied to Kalman filtering in Equation (23), which is a linear-dynamical system (LDS) evolves with the transition matrix computed for each time step. We briefly mentioned this in the next sentence: "In theory, we could jump to any directly in the Kalman update in Equation (23), but that would increase the uncertainty as if there was no data before, changing the dynamics."

Same paragraph: “Instead, we also compute all intermediate steps, which mirrors the more realistic setting where we also observe data at 1, 2, . . . , t − 1.” – Is the initialization of UnLSSL therefore data dependent?

No, it just means that we compute the parameters for the case where we observe data at each integer time step. Otherwise, the computed parameters would be computed as if we had not observed any data for longer stretches of time and therefore increase the uncertainty estimate and smoothing.

How were the noise experiments conducted? What type of noise was added? At what SNR?

For each audio signal, we compute the standard deviation of the signal and add a random noise with times the computed standard deviation. We used 10 different values: 0.0, 1e-7, 1e-6, 3.16e-6, 1e-5, 1.77e-5, 3.16e-5, 5.62e-5, 1e-4. We observed values larger than 1e-4 repress most of the audio signals in SC10 and leaves nothing but noise for this dataset.

Since we do not add the noise based on amplitudes, and instead based on the standard deviation, we do not have a constant SNR value. For example, an audio with constant signal would never be effected by noise even when is non-zero as the standard deviation of the signal itself is 0.

We run the LSSL model with HiPPO and UnHiPPO initialization on the same data and report averages of three different seeds. The complete configuration for the noise experiments is available in the supplementary material in config/experiment/sc-raw-noise.yaml.

Thank you for your detailed review.

Time-dependence of matrices

It is correct that SSMs learn $A$ and $B$ and discretize them on the fly. However, at least for LSSL, the time step at which the learned matrices are discretized is fixed per feature to make the model adapt to multiple timescales and does not actually vary with time. Therefore, it is factually equivalent that we learn $A_k$ and $B_k$ directly.

Experiments

We did not evaluate on LRA, because its synthetic tasks have discrete input data, which does not fit the assumptions of our derivation. In there, we assume that the noise is Gaussian and therefore continuous. An adaptation for discrete data and discrete noise would be a nontrivial adaptation.

Typo and related work

Thank you for your careful reading. We have added the missing "by" to clarify that the GELU nonlinearity comes first and included the two works you mentioned in our related work section.

Compatibility with efficient parametrizations

As we described under Limitations (Section 8), we were unfortunately not able to derive a similarly structured representation of our initialization, because the pseudo-inverse and the Kalman filter equations pose a significant challenge, which we elaborate in the following.

$\\mathbf{A}_{\\mathrm{R}}$ is not a normal matrix, therefore, there is no trivial diagonalization of it. Previously, Gu et. al. (S4) used correction matrices to obtain skew-symmetric matrices which are in turn a special case of normal matrices. Here, though, finding a correction matrix is not trivial due to pseudo-inverse. Considering the definition of pseudo-inverse for matrices constituting of linearly independent columns,

\\mathbf{M}^{\\dagger} = (\\mathbf{M}^{\\ast}\\mathbf{M})^{-1}\\mathbf{M}^{\\ast}\

where $(\\cdot)^{\\ast}$ denotes the conjugate transpose that is equivalent to transpose for a real matrix. Since $\\mathbf{A}_{\\mathrm{R}}$ is real, we can safely use,

\\mathbf{M}^{\\dagger} = (\\mathbf{M}^{\\mathsf{T}} \\mathbf{M})^{\-1}\\mathbf{M}^{\\mathsf{T}}\

Then, pseudo-inverse of (19) becomes

$

\begin{pmatrix}\mathbf{I}\\ \mathbf{B}^{\mathsf{T}}_{\text{H}}\\ \mathbf{Q}^{\mathsf{T}}\end{pmatrix}^{\dagger} = \begin{pmatrix}\begin{pmatrix}\mathbf{I}\\ \mathbf{B}^{\mathsf{T}}_{\text{H}}\\ \mathbf{Q}^{\mathsf{T}}\end{pmatrix} \begin{pmatrix}\mathbf{I} \\ \mathbf{B}^{\mathsf{T}}_{\text{H}}\\ \mathbf{Q}^{\mathsf{T}}\end{pmatrix}\end{pmatrix}^{-1} \begin{pmatrix}\mathbf{I}\\ \mathbf{B}^{\mathsf{T}}_{\text{H}}\\ \mathbf{Q}^{\mathsf{T}}\end{pmatrix}^{\mathsf{T}}\nonumber
$

and $\\mathbf{A}_{\\mathrm{R}}$ can be explicitly written as

$

\mathbf{A}_{\mathrm{R}} = \begin{pmatrix}\underbrace{{\begin{pmatrix}\mathbf{I} \\ \mathbf{B}_{\text{H}}^{\mathsf{T}} \\ \mathbf{Q}^{\mathsf{T}}\end{pmatrix}}^{\mathsf{T}}{\begin{pmatrix}\mathbf{I} \\ \mathbf{B}_{\text{H}}^{\mathsf{T}} \\ \mathbf{Q}^{\mathsf{T}}\end{pmatrix}}}_{\text{First part}} \end{pmatrix}^{-1} \underbrace{{\begin{pmatrix}\mathbf{I} \\ \mathbf{B}_{\text{H}}^{\mathsf{T}} \\ \mathbf{Q}^{\mathsf{T}}\end{pmatrix}}^{\mathsf{T}} \begin{pmatrix}\mathbf{A}_{\text{H}}^{\mathsf{T}} - \mathbf{I} \\ 2\mathbf{Q}^{\mathsf{T}} \\ \mathbf{0}\end{pmatrix}}_{\text{Second part}}
$

First part gives us a symmetric and square matrix. Therefore, we know that inverse of this will be also symmetric, and this part would not be a problem for diagonalization if it were to stand by itself, or alongside another symmetric matrix.

Second part does not result in a symmetric, skew-symmetric or in general a normal matrix. Therefore, the overall result needs a correction matrix to become diagonalizable matrix.

However, hand calculation requires taking the inverse of the first part. Even after that, multiplication with a non-normal matrix needs to be computed, together which makes it too complicated to reveal the structure of the matrix, so that it can be manipulated to have a normal form.