WaLRUS: Wavelets for Long range Representation Using State Space Methods

Dear Reviewer,

Thank you for your feedback. We would like to address some of your concerns about this work in order to strengthen and improve our paper. Brief summaries of your notes are in bold, with our responses below.

“There seems to be no code supplement for reproducibility purposes.”

A link to the code supplement was provided in our paper twice. Please see the footnote on page 1, and the NeurIPS checklist, item 5.

“...it is not shown how to calculate the transition matrices A and B using the wavelet basis.”

Equations 3 and 4 give the construction of A and B explicitly. These equations and their derivations can be found in reference [17], as noted in the text. There is also an algorithmic implementation in our code.

The construction of wavelet frames is well-covered in the literature (see references [18],[20],[29],[30]). Figure 2 contains an illustration of a partial Daubechies wavelet frame, providing readers with key information (scales, shifts, and function shape) required to understand the contribution of this paper, without needing full knowledge of the vast literature on wavelet frame construction.

To facilitate easy access to key information for our readers, we have also added a section in the Appendix of our revised manuscript with details from [17] on the derivations of A and B.

“It is claimed that wavelet bases always yield diagonalizable state-space representations, but this is not justified”

We did not claim that wavelet bases “always yield” diagonalizable SSMs; we were careful to note that wavelet-based SSMs were found to be empirically stably diagonalizable (line 47-51) and observed to be diagonalizable in our experiments (line 137). Prior work has not proved any theory regarding which bases/frames or measures produce stably diagonalizable $A$ matrices, and have thus far only reported empirical observations, as we have done.

Critically, the only significant previous work on this subject (HiPPO) found that even at fairly low rank, Legendre-based matrices are not stably diagonalizable. This problem was a major roadblock in the development of S4 and ultimately Mamba, and provided the motivation for multiple subsequent follow-up papers on how to make diagonal approximations to the Legendre-based $A$ matrix (see refs [12-14]). In fact, if one examines the code for the state-of-the-art Mamba model (https://github.com/state-spaces/mamba), the SSM layers are initialized with one of these approximations, and not the true $A$ matrix. Thus, we believe it is valuable to note our empirical findings of diagonalizability, even in the absence of a full theoretical proof.

To make very clear the limits of our claims, we have made the following wording adjustments in our revised manuscript:

Line 47-51, Section 1: “Empirical results demonstrate … WaLRUS enjoys diagonalizability, which is the key enabler of efficient convolution-based implementations and parallel computation [13, 14].”

Change to → “Empirical results demonstrate … WaLRUS has been consistently experimentally observed to be stably diagonalizable, which is the key enabler of efficient convolution-based implementations and parallel computation [13, 14].”

Line 137, Section 3.2: “We observe that both scaled and translated WaLRUS are stably diagonalizable.”

Change to → “We observed that each of the scaled and translated WaLRUS SSMs we implemented, regardless of dimension, were stably diagonalizable. Further research is required to determine whether Daubechies wavelets will always yield diagonalizable SSMs.”

While a complete theoretical proof of diagonalizability is outside the scope of this particular work, we agree that it would be a significant contribution to the literature on this topic. We hope that our paper inspires members of our community to join us in exploring this problem.

“Only state-space methods are compared with each other, while competing methods (e.g., RNNs, Markov methods) are left out.”

RNNs and Markov methods are not competing methods with online basis projection. Our paper proposes an alternative basis for state-of-the-art function approximation SSMs that play a key role in some of the fastest-growing models in ML today; it does not propose a time-series predictor.

HiPPO / SaFARi / WaLRUS only store previous information by iteratively (and deterministically) projecting a signal onto a basis; probabilistic methods such Markov methods are designed for predictive tasks. Most importantly, neither RNNs nor Markov methods produce a model compatible with use in S4, Mamba, and Mamba’s variants, a driving motivation for this work. Thus, the only baseline for comparison is existing SSMs for function approximation (namely HiPPO).

“it is unclear how the minimization on line 70 yields a state vector”

The results in line 70 were a key finding of the original HiPPO paper [9]; we do not reproduce it in our paper, only noting the results pertinent to our work. To ensure that this paragraph does not cause confusion for readers unfamiliar with HiPPO, we have made the following wording change in our revised manuscript, removing the reference to the minimization:

Line 70: “To resolve these issues, the HiPPO framework [9] casts online function approximation as a continuous projection of the input u(t) onto a linear combination of the given basis functions G. At every time T, it solves ming (T )∈G ∥uT − g (T) 70 (t)∥µ , producing produces a compressed state vector c(T) that satisfies the update rule:”

We thank the reviewer again for their time and the opportunity to clarify and improve our paper.

Dear Reviewer,

Thank you very much for your thorough and insightful feedback. We would like to address some of your concerns about this work in order to strengthen our paper. Brief summaries of your feedback are in bold, with our responses below.

“The authors switch a part in an existing framework. While this is shown to be better, the overall novelty is not of the highest quality… the shown improvements in MSE are not that high. The relevance of those improvements to downstream tasks remain uncertain.”

We have framed our work as a drop-in replacement for SSM-based ML models as motivation and context, but this work also stands alone as a lightweight online function approximation tool. More importantly, it is the very first work offering a replacement for the Legendre-based SSM that is taken for granted when utilizing off-the-shelf Mamba/S4 implementations. To date, no other work has ever gone back to first principles to understand the functionality of the core machinery running these models, let alone fundamentally alter it.

Based on results from prior and current work, we believe that using better-performing function-approximating SSMs in neural networks should improve performance in downstream tasks; however, we first must show how and why these SSMs perform better on their own. The focus of this paper is thus on developing a core building block that will enable those downstream applications, and hopefully be of use and interest to members of the community.

Our revised manuscript includes a new figure--a simplified block diagram illustrating the relationship between our work, HiPPO, SaFARi, and Mamba. We have also included some additional text in Section 2 that contextualizes our work in the space of Mamba and other such models:

"Some recent work [R1,R2] has conceptually connected the use of wavelets and SSM-based models (namely Mamba). These efforts are fundamentally distinct from ours in that they perform a multi-resolution analysis on the input to the model only. No change is made to the standard Mamba SSM layer; it is still using Legendre-polynomial based SSMs.

This work, on the other hand, is the first to challenge the ubiquity of the Legendre-based SSM, and present alternative wavelet-based machinery for the core of powerful models like Mamba. WaLRUS could be used as a drop-in replacement for any existing SSM-based framework. However, before simply substituting a part in a larger system, we must first justify how and why a different SSM can improve performance. This paper presents a tool that stands alone as an online function approximator, and also provides a foundational building block for future integration in SSM-based models.”

“Wavelet frames come with more parameters, e.g. scale, shift N_eff. It is totally unclear how to set them. There are some justifications for the set parameters, but it is still unclear how to set them properly. This is however mentioned in the limitations.”

Scales for Daubechies wavelets are always dyadic, scaling by a factor of 2; we have updated the equation in Section 2.1 to make this clearer:

$$W(j,k) = \int_{-\infty}^\infty f(t) \psi^*_{j,k}(t) \, dt, \quad \psi_{j,k}(t) = \frac{1}{\sqrt{2^{-j}}} \psi\left(\frac{t-k}{2^{-j}}\right)$$

The choice of shift amount introduces some interesting complexity that we agree could be given deeper treatment, and we have added the following text to Section 3:

“To construct the frame, we use the usual dyadic scaling for multiresolution analysis; that is, scaling the mother wavelets by a factor of two at each level. For each scale, different shifts along the x-axis are introduced. Compressive wavelet frames are truncated versions of wavelet frames that contain only a few of the coarser scales, and introduce overlapping shifts to keep the expressivity and satisfy the frame condition (See Mallat, [18]). The interplay between the retained scales and the minimum required overlap to maintain the expressivity is extensively studied in the wavelet literature [18,20,30]. If there is excess overlap in shifts, the wavelet frame becomes redundant, and redundancy has advantages in expressivity and robustness to noise.

The exact relationship between the amount of overlap and the resulting dimensionality of the A matrix depends not only on the overlap of the shifts, but also on the type (and order) of chosen wavelet, and number of scales. Determining the “optimal” overlap or N_eff is application-specific, and an area for future research.”

“A.1.3 lists all the chosen parameters, which vary from experiment to experiment. That seems like hyperparameter optimization and making the significance of the result somewhat questionable.”

Some parameter choices are heuristic and some are codependent, but none were chosen to be optimized, or to otherwise unfairly advantage one method over another. We acknowledge that the varying parameters as presented in Table 3 could reasonably raise this concern for readers, however, and so we have added additional text to clearly explain how these parameters were chosen. We have also simplified our table to remove the wavelet type (D22, which is the same for every experiment in this work and is already described in the same appendix), and the “scale max” variable, which does not impact the rank of the frame, as explained in the added text below.

“Scale Min and $N_{\rm{eff}}$: The minimum scale sets the smallest feature of the signal that can be represented by the frame. This parameter should be chosen based on knowledge about the signal of interest and its component frequencies, as one would for a filter bank. Note that the size of the smallest feature is relative to the length of the signal under consideration, so this value may differ under scaling and translating measures.

Scale min also controls the effective rank, $N_{\rm{eff}}$. Each additional lower scale introduces a factor of two in the effective rank of the frame, owing to the additional shifted elements in each scale. To illustrate, Fig. 2 shows two scales, where there are three father wavelets ($\phi_0$) and three coarse-scale mother wavelets ($\psi_1$). The next scale introduces six scaled and shifted mother wavelets ($\psi_2$), the next would include 12, and so on. Table 3 also shows this pattern, with a minimum scale of 0 corresponding to $N_{\rm{eff}}$ of $2^6$, minimum scale of $−1$ corresponding to $N_{\rm{eff}}$ of $2^7$, and so on, with variance resulting from the shift combined with truncation of frame elements during construction.

Our code includes another variable, scale max. Since smaller scales can also combine to represent larger scales, scale max in fact has no impact on $N_{\rm{eff}}$. (See Mallat [18] for further details on this topic.) Adding coarser scales can help improve numerical accuracy in the calculation of $A$, however. We do not include scale max in Table 3, but we provide it in our code with each experiment for reproducibility."

To support this additional text, we also have generated some quantitative results. Fig. 10 (in the appendix of our revised paper) takes an example implementation of WaLRUS and varies scale max, and we observe that the size of $N_{\rm{eff}}$ remains constant. It is also easily shown that varying scale max results in the same diagonalized $A$ matrix.

“What is the effect of varying the effective dimension N_eff? Could you plot MSE over N_eff?”

We agree that analyzing the impact of varying $N_{\rm{eff}}$​ on MSE would be insightful. However, due to limited time, we cannot include this experiment in the current revision, and plan to pursue it in future work.

“The critique in section 3 on other frames seems to focus on the Legendre translation (LegT) frame, other frames occupy less room. Can you balance this more?”

We have added the following text in Section 3 of our revised manuscript after the discussion of the translated Legendre kernel:

"The translated Fourier (FouT) kernel primarily suffers from the opposite problem: substantial nonzero elements outside the kernel window indicate that LegT struggles to effectively "forget'' historical input values. Thus contributions from input signals outside the sliding window appear as representation errors. LegT also has this problem, to a lesser extent--see area A of Fig. 3."

We keep the primary focus on Legendre-based kernels, as this is the most ubiquitous (even if hidden!) SSM used today in ML architectures.

“The metrics for the "isolated spikes" task are only described with an image. Could you provide some formulas as well?”

Yes; we have added a section in the appendix of our revised manuscript with these formulas.

We thank the reviewer again for their thoughtful questions, which spurred interesting discussions among the authors and provided concrete ways to better communicate the content and context of this work. We believe our updated manuscript is significantly improved, and look forward to any further feedback you would provide.

---

References:

[R1] Zou, W., Gao, H., Yang, W., & Liu, T. (2024, October). Wave-mamba: Wavelet state space model for ultra-high-definition low-light image enhancement. In Proceedings of the 32nd ACM International Conference on Multimedia (pp. 1534-1543).

[R2] Zhang, T., Zhu, Y., Zhao, J., Cui, G., & Zheng, Y. (2025). Exploring state space model in wavelet domain: An infrared and visible image fusion network via wavelet transform and state space model. arXiv preprint arXiv:2503.18378.

Dear Reviewer,

Thank you for your thorough and insightful feedback. We would like to address some of your concerns about this work in order to strengthen our paper and better tailor it to a NeurIPS audience. Brief summaries of your feedback are in bold, with our responses below.

“It would be nice to provide a bit more exposure / explanation to Eq 3,4 which use notation that not everybody in this ML venue may be familiar with.”

We agree that since this is very recent work that readers may not be familiar with, it will be helpful to provide deeper insight into these equations. We have written a new appendix section which provides a quick and convenient reference for readers on the derivations of Eqns 3 and 4 following reference [17]. This section starts from a simple projection of a signal onto a basis via inner product, and develops the full SSM models step by step with text narrative to build intuition along with the mathematical framework.

“The narrative of the abstract/intro is: There is HiPPO --> SaFARi --> WaLRUS (based on SaFARI). But the experiments compare WaLRUS only against two HIPPO variants. Why not compare against the direct baseline SaFARi?”

We have added clarifying text to Section 1 to emphasize the point that SaFARi is a method to generate models, whereas WaLRUS is a model instance:

“In this paper, we leverage the SaFARi method with wavelet frames to introduce a new SSM instance, WaLRUS (Wavelets for Long-range Representation Using SSMs).”

We have also created a block diagram to illustrate the relationships between WaLRUS, HiPPO, SaFARi, and Mamba, which also helps to highlight the unique contributions of this paper and the distinction from other works.

“How reusable is the proposed Walrus code? I.e. from reading the paper I'm not sure how easy it will be to try this approach out as drop-in replacement in existing SSM frameworks.”

WaLRUS is indeed a drop-in replacement for existing SSM frameworks using HiPPO; the new figure in our revision explains this. Our code generates A and B matrices of size NxN and Nx1 respectively, which can be directly substituted into any code that uses SSM layers made from A and B matrices (code supplement provided in footnote on p.1).

“There seem to be few works [1], [2], that do Wavelets with SSM… it begs the question whether Wavelet + SSM is novel - or what are the gaps that this works fills exactly?”

We did not cite extensively from the vast and growing literature on Mamba, since it is outside the scope of our paper, except as motivation and contextualization. Yet, given the overlap in the titles of these papers and the focus of our work, we expect that other readers will have the same question. Therefore, we have include these citations in Section 2 with the following text:

"Some recent work [32, 33] has conceptually connected the use of wavelets and SSM-based models (namely Mamba). These efforts are fundamentally distinct from ours in that they perform a multi-resolution analysis on the input to the model only. No change is made to the standard Mamba SSM layer; it is still using Legendre-polynomial based SSMs.

This work, on the other hand, is the first to challenge the ubiquity of the Legendre-based SSM, and present alternative wavelet-based machinery for the core of powerful models like Mamba. WaLRUS could be used as a drop-in replacement for any existing SSM-based framework. However, before simply substituting a part in a larger system, we must first justify how and why a different SSM can improve performance. This paper presents a tool that stands alone as an online function approximator, and also provides a foundational building block for future integration in SSM-based models.”

“Do the current benchmarks really demonstrate WaLRUS’s benefits over existing SSMs (and maybe even Transformers)? I believe an experiment on more community-relevant settings a small-scale long-context language-model or long biological sequences could provide a clearer picture.” “The experimental evaluations are on tasks that I personally do not find super relevant to showcase the benefit of a new SSM variant.”

WaLRUS, like HiPPO, does not learn any parameters; it is a purely signal-processing approach that does not compare directly to learned models like transformers. It holds the promise of building better machinery for ML models like Mamba, but it also stands on its own as a function approximation technique. Our goal in choosing experimental cases was to highlight the power of WaLRUS in function approximation as a first step in developing a new tool, and as such, the only appropriate comparison is to HiPPO. We again use our new diagram to emphasize this focus.

The use cases you suggested – language models and biological sequences – are excellent candidates for future work, e.g. by comparing a WaLRUS-based Mamba to a HiPPO-based Mamba. In fact, we hope that this work will inspire the community to apply WaLRUS in such scenarios by building on this foundational work.

Formatting of Table 2: We agree that without vertical lines it is much nicer; we have removed these in our most recent revision.

We believe our updated manuscript is significantly improved and better contextualizes the work. We thank the reviewer again for their time and for the valuable feedback on our paper, and look forward to any additional comments you would provide.