Flash Inference: Near Linear Time Inference for Long Convolution Sequence Models and Beyond

Thank you for the comments! We tried to address them in the first revision of the paper and below:

Regarding weaknesses:

• We use a different tiling from that of van der Hoeven to exploit data-independent filters. We have a significant number of systems improvements to make the method practical, that we enumerate and emphasize in several parts in the paper (in the intro, as well as the improvements justification section). Among these, the across-layer parallelization stands out as an optimization that is specific to such layered architectures. We also show how to halven the memory usage, how to exploit DFT precomputations, dynamically choosing the best way to deal with a tile and so on - generally doing whatever it took to efficiently implement our method on a GPU. Finally we generalize the method to formulate a framework that can help guide the design of new architectures.
• We will touch on this in our upcoming extended discussion on the context/literature review. However, since our method works for any convolutions, any SSM could benefit from it - it is just that if they are low-dimensional enough, the recurrent view could be faster, but our method can be applied as well. Asymptotically, assuming SISO (single-input-single-output), one would benefit from our method if the state dimension exceeds $\Omega(\log(L)^2)$. Besides SSMs, we hope the method has the potential to inspire future architectures (potentially transformers’ adaptations) and this is precisely why we propose the general framework mentioned in Section 3.4 and Appendix “The Flash Inference Framework”.
• In our new revision we added the complete sweep mentioned in the Experiments section explicitly in the paper, together with some of the associated plots.

Regarding the questions:

• We did so above.
• Following the integration of CUDA Graphs, we dropped the non-mixer times by a significant constant (since that part was dominated by CPU overhead). Following this revision, we now get very significant end-to-end improvements as well. Furthermore, even in the initial case we dropped the mixer part from the same order of magnitude with the MLPs to negligible - this then points to MLPs being a part to tackle next (say via structure matrix approaches). Regarding the improvements potentially becoming marginal, on the actual hyena architecture, we get at least $2.8\times$ improvements in the new experiments on every single hyperparameter combination (and up to almost $8\times$).
• The reductions of Hybrid vs FFT are shown (numerically) for all experiments in the tables of Appendix Further Experimental Details, they are sometimes marginal, but sometimes get up to 10% extra performance. Regarding the non-mixer time being smaller for the eager/lazy counterparts in Figure 3c, this has been fixed via the usage of CUDA Graphs, but the initial cause was already covered in the caption as well as in Section 4.3, namely better mixers exposing CPU overhead. As mentioned above and in the general rebuttal, we got rid of this issue. We have added a discussion on this in Appendix D.2 (CPU overhead: the silent performance killer).

Thank you for the insightful comments! In our first revision we tried to cover all of them except for weakness 2 and question 3 that we are currently working on.

Regarding the weaknesses:

• W1:
  • At line 163, we changed to hadamard notation to not make the channels explicit. Since the channels appear very sparingly, we decided to use the superscript for layers instead in hope that this makes the formulas more readable (and let channels/features be subscripts).
  • We changed it to denote it being a definition.
  • We updated the reference.
• W2: We appreciate your points and will update our paper in the next revision to integrate a deeper and more nuanced discussion as requested. We wanted to make this first update to respond to the smaller points and update on our improved experiments but will be back to addressing this. Regarding the “popular misconception”, we were referencing the ML community where the bias mostly stems from these models having been proposed following low-dimensional LTI SSMs - however, we’ll change to something less subjective, we simply wanted to emphasize the fact that LCSMs' equivalent LTI SSM representation need not be low-dimensional (and thus, that general LCSM' inference is not settled through the recurrent approach).
• W3:
  • We considered it is common practice to not denote the base of the logarithm since the asymptotics of different bases stay the same: that is, $O(log_2(L))=O(log_{100}(L))$. We do not try to make any clear claims about the constant involved since that is a function of FFT implementation and more.
  • We added to Appendix “Discussion on Memory” a deeper discussion on the relationship between across-layer parallelization and memory bandwidth.
  • We did not intend to propose a new architecture ourselves, but with regards to whether we need data-dependency, we referenced two works (Zoology and Orchid) in the conclusion that show the potential of such methods - however they did not succeed yet in finding a causal alternative. We added a mention to these in the main body at line 235 as well.

Regarding the questions:

• In our revised version, we added experiments on higher context lengths (up to 128K and even 256K). There is nothing special about 32K, we only chose that plot since the CPU overhead was more stable across methods, but following the integration of CUDA Graphs, all the experiments have this property. As is expected, the higher the context length, the higher the mixer speedup. Any specific issues for certain context lengths would be due the FFT implementation (which past a certain point will struggle to fit the relevant sequences in the SRAM and thus incur a slowdown), but one would still expect this approach to be better than lazy/eager.
• We have added this to Appendix “Framework Adaptations”, but in short the method can be applied to the relevant subset of layers that fit the framework. For example, one could use the same approach for Striped Hyena - even the across layer parallelization persists, but it can only happen on the relevant subset.
• We agree this is a useful experiment to have and while it is difficult to measure accuracy tradeoffs since this is architecture and data-dependent, it certainly would be good to know from what SSM dimension onwards our method works better - theory would predict $\log^2(L)$ assuming SISO, but the constant could vary a lot. We are trying our best to get this experiment up in due time, but will mention at least the theoretical asymptotic. Regarding the linear growth of the cache, this seems fundamentally needed if one wants to use absolutely any filter length, but of course it is true that we could settle for particular structures.
• We are actively looking into this. Interestingly, this would be easy to integrate if one assumes the filter to be sparse in the first few positions (and one could use some different layer for local behavior) - simply by dropping dealing with the small tiles. This would have the same theoretical complexity but probably improve the constant. Thank you for the suggestion!

We just updated to our second revision where we added an extended contextual discussion.

Specifically, regarding the second weakness:

• We dropped the strong statement and replaced it by just mentioning that we fundamentally change the class of models. We discuss in Appendix A the advantages of having a fast inference framework that does not rely on particular biases and give a more nuanced discussion on when it is desirable to choose our method compared to others that exploit certain structural properties (including SSMs).
• We cover prefilling vs autoregressive stage in Appendix A.3. The only thing different from transformers is the eager nature of the prefill which lets us drop the first P activations (where P is the length of the prompt). We also mentioned how our improvements are most relevant when the prompt size does not dominate the number of generated tokens.
• We dropped the claim altogether and rewrote that paragraph, redirecting to Appendix A for more background. We only state that there is a bijection between SSMs and LCSMs and that the equivalent SSM to our filter of interest could have too big of a dimension for the recurrent view to be efficient, while Flash Inference has a filter-independent cost.

We hope this addressed your concerns!

Thank you for acknowledging our contribution and for the comments! Our responses are below and we have revised our paper following them.

Regarding the weaknesses:

• We theoretically proved that the method applies to any LCSM (and LTI SSM for that matter). This includes by definition CKConv, FlexConv, H3, etc. Our synthetic ones already cover the general form of LCSMs: whether you place an MLP or some other kind of block, the mixer times stay the same and that is where our improvement comes from. We performed the Hyena experiments to prove this synthetic-to-complete-architecture transition, especially since Hyena is more involved than its counterparts. Regarding real-word tasks, could you please explain what sort of performance you refer to? Our method performs exact inference and the time does not depend on the weights, so the quality of a model stays the same. We also do not see how experimenting with any other LCSM architecture would add much - we only optimize the convolution part which we assume to be the same regardless of architecture (we even share the tile-handling code between our synthetic and hyena experiments), so we kindly ask you to let us know what concerns you have about the current setup.
• We added an extra note at line 214 mentioning that data-dependent filters incur a factor of 2 (this was done implicitly at line 233 anyway). We covered at line 726 (Appendix C) how data-dependent filters simply require a different tiling that at each step will perform 2 FFT calls of the same size instead of one. Our whole theoretical analysis extends and given the very simple nature of the modification, one can easily expect the implementation to take at most twice the amount of time (it could be even faster than that for tiles that didn’t use to be memory or compute bound).
• We added a discussion in the Appendix (Framework Adaptations) on this. In short, we make the absolute minimal assumptions needed for the model to well define autoregressive generation. In particular, one can accommodate non-causal encoding of a prompt for example or even chatbot-like interactions.

Regarding the questions:

• We have mentioned in Section 3.4 both why transformers cannot benefit from Flash Inference as well as a high-level description of what class of architectures our method can be applied to. The class is thoroughly defined in Appendix B. We could not identify any existing architectures beyond LCSMs that fit the description, but this is precisely the purpose of the generalization: to encourage people to try out such architectures - as suggested as a future direction in the second point of conclusions.
• As mentioned above, we will incur at most a factor of $2$ efficiency cost practically (on the mixer side).
• We could not find any well-defined way to compare inference for different architectures: is it for the same number of parameters or for the same compute training budget? Or for the same embedding dimension? LCSM’s parameters can be distributed very differently from those of transformers (since the attention layer has the Q, K, V matrices parameterize it whereas, for the same embedding dimension, the parameterization of the SISO convolutional filters could vary wildly in LCSMs), so as far as a fair direct comparison goes, we can only look at the theoretical asymptotics in $L$ (where attention takes $O(L^2)$ and we have $O(L\log^2(L))$).
• We covered this above.

Thank you for the thoughtful comments!

We are happy to try and integrate them, but before doing so we have a few questions to make sure we address the right issues:

• Regarding the first weakness, we do provide Algorithm 2 which is the complete algorithm for an end-to-end LCSM model as well as Algorithm 1 which deals solely with one convolution layer. We could, if you find this to address your problem, add a python implementation for one layer (it would be mostly the same as Algorithm 1 but with an extra MLP and the $\tau$ function implementation based on FFT). Otherwise, could you please pinpoint which parts of Algorithm 2 are too abstract? Could you also please detail what design decisions you have in mind? We are not aware of any implicit decisions, but are happy to add more details about any part.
• For the second weakness
  • The function $\tau$ has first been defined at lines 261-262 (Eq. 3) for the case of relaxed polynomial interpolation and then overwritten for the general case at line 338. We intentionally use the same symbol since it has the same function, just across more dimensions, so that it is easier to follow the transition from Algorithm 1 to Algorithm 2. It is true that the first definition is as part of Lemma 1's statement: we are happy to define it before we state the Lemma to avoid any confusions.
  • We will add in the Appendix the other 3 implementations of $\tau$, but regarding how we choose between them: this is covered at lines 478-479. Specifically, we profile each implementation for each batch dimension, embedding dimensions, number of layers and tile size and automatically choose the best one for each iteration. There was indeed a typo at the end of the sentence: we should've referenced Figure 3a which shows the runtimes for the four implementations for varying tile sizes. We can provide the specific implementation choice for each (B, D, M, U) quadruple in the Appendix if that helps, but generally the pattern is as showed in Figure 3a: the quadratic implementations are better for small tiles and then the FFT ones dominate with the torch vs Flash implementations sometimes flipping.

Regarding the first question, thank you for the suggestion, we will add it. As for the second question, what we meant is that van der Hoeven's approach could be used as well, but since $\rho$ is known ahead of time we can do better. We hope the following restatement is clearer:

Whereas one could simply use the more general approach of (van der Hoeven, 1997), we can further take advantage of ρ being known ahead of time to get a twice-faster and slightly simpler algorithm

The reason for the slightly more abstract view is simply to cover all LCSMs at once. We are happy to address any specific design choice you may have in mind and implement any actionable advice.

I think there is a fundamental confusion here about definitions. Do we agree or disagree: by your definition of LCSM (section 2.3), all SSMs (henceforth assumed to be LTI and SISO unless otherwise qualified) are LCSMs?

the only SSM that makes use of the convolutional view (because it has a meaningfully large state dimension) is S4 - all the others exploit a small state dimension which LCSMs do not assume.

First, all SSMs mentioned in my original review ("LSSL, S4, DSS/S4D, H3, MEGA/Megalodon, and many more") do use the convolutional view. This has nothing to do with state size.

but as stated in paragraph 2, we acknowledge that any convolution has an equivalent LTI SSM - it’s the dimension of that state we emphasize as being the core difference

Second, my claim is not about any convolution having an equivalent LTI SSM, but the converse: all SSMs are by definition LCSMs.

These claims have nothing to do with the state size of the SSMs, for now. Let's get on the same page about this before I respond to the rest.

We thank the reviewer for the thorough comments and for acknowledging our technical contribution!

Regarding the weaknesses:

• While we are happy to add more context (and thank you for the pointers, have missed FlexConv), the only SSM that makes use of the convolutional view (because it has a meaningfully large state dimension) is S4 - all the others exploit a small state dimension which LCSMs do not assume. We will add a section in the appendix on the connection between LCSMs and SSMs, but as stated in paragraph 2, we acknowledge that any convolution has an equivalent LTI SSM - it’s the dimension of that state we emphasize as being the core difference. The reason we did not delve deeper into this is simply because the purpose of the paper is not to argue for the use of LCSMs, but merely to provide an efficient way of serving them - an open problem that was also approached in Laughing Hyena Distillery.

• Regarding the anchoring around Hyena, we will drop (or broaden) the references at lines 17, 43, 213 and 530. Thank you for the pointer - we thought Hyena is the most well known of the LCSMs and wanted to make explicit which parts apply to it as well. Moreover, the closest work to ours in terms of problem setup is Laughing Hyena Distillery. Do you mind elaborating what performance you would find meaningful to evaluate? Since we do exact inference, we only claim a certain speed-up: anything else related to the statistical quality of the model stays unchanged. Regarding why we tried only Hyena: is there any architecture you believe running extra experiments on would add any value? The reasoning was that the synthetic setting should cover all LCSMs and Hyena is just an extra example. Your last point about the method being generic is true and our synthetic setting, together with Algorithm 2 are not specific to Hyena in any way.

• Firstly, Fig.1 of the linked paper does not seem to have any data on H3/Hyena. Secondly, and more importantly, it is very challenging to claim complete superiority of an architecture over another, especially in the opposite direction of representational power, simply because one cannot consider all the ways that module could be used: performance could vary a lot depending on use case and integration. Our method speeds up a long convolution layer, which is generic enough to be considered a useful building block for yet-unexplored scenarios. It could be that H3>Hyena by themselves and hybrids of attention with H3 are worse than hybrids of attention with Hyena. To give a precise example, EVO (which is based on Striped Hyena) is the current state-of-the-art in DNA modeling. Moreover, one could always change the biases of the convolutional filters: the main issue is that of training stability, which is itself an open problem. Our method serves as a tool readily usable for any such method. We believe it is useful for the community to be aware of a fast inference alternative for them to not constrain their research to low-dimensional SSMs because of computational considerations.

Regarding the questions:

• EVO (based on striped hyena) is an example of a long convolution-based model that is more performant than the equivalent SSMs (EVO tried mamba which failed due to numerical instability). We unfortunately cannot easily test how much of the accuracy is lost when distilling in due time, but the model is trained in the convolutional, not SSM view. S4, although primarily viewed as an SSM, is still SOTA on the Image section of the Long Range Arena tasks and requires a dimension equal to the sequence length.
• Orchid and Zoology specifically show the representational advantages of data-dependent convolutions to synthetic tasks of interest - both of which suggest non-causal architectures which outperform SSM-based counterparts. It is an open problem to find a good causal architectural choice that is stable to training, but it is clear that the added representation power closes a gap.

Regarding the broader positioning, we believe it is important to not overlook one main contribution: that of defining exact properties that allow for similar improvements to be applied to other architectures. We consider this is a good reason for the paper to be published at a conference such as ICLR since people can take these considerations into account when designing the next generations of architectures. While this is touched just briefly in the main paper due to limited space, we have a complete analysis in the appendix. Furthermore, we hope that the technical part of our method inspires people to think differently about similar problems that may arise in their practical research. Finally, we would also appreciate if you detailed what you meant through “real models”, since we do have a realistic Hyena implementation which is a real model.

Some of my other points are about the utility of such models. As of right now, there are no results that suggest that there exist general LCSMs that outperform efficient LCSMs.

6. In all benchmarks where controlled comparisons exist (e.g. not including the EVO model, where no third party reproductions exist or ablations between the Hyena parameterization vs other LCSM parameterizations), the performance of non-SSM LCSMs and SSMs are essentially the same.

7. E.g. as the authors note "S4, although primarily viewed as an SSM, is still SOTA on the Image section of the Long Range Arena tasks and requires a dimension equal to the sequence length." Actually S4 and variants (MEGA, etc.) all use a low dimensional SSM, and the majority of methods near SOTA on LRA are based on low-dimensional SSMs or equivalent. I will also point out that the authors' phrasing "although primarily viewed as an SSM" once again implies a false dichotomy between SSMs and LCSMs that is being perpetuated. SSMs are simply a class of LCSM with additional properties.

8. There is in fact a technical reason for this, which was touched on in the original review. As mentioned in point 4 above, all LCSMs have restricted expressivity. This is essentially tight, in the sense that any LCSM with $P$ parameters requires $O(PL)$ time to construct its convolution filter, so there is a fundamental tradeoff in the expressivity <-> training time of LCSMs; both implicit convolution models like CKConv/Hyena as well as SSMs attain this bound. So there is no intrinsic expressivity advantage to using a non-SSM LCSM over an SSM (unlike what is often implied by the paper; for example, pointing out that general convolutions cannot be represented by a low-dimensional SSM is misleading, as they cannot be represented by any other finite-parameter class of LCSM either). Thus the question to ask is not about expressivity but about whether a particular parameterization carries a helpful inductive bias. However, the best non-efficient LCSMs such as Hyena and SGConv actually bake in priors like exponential decay which intentionally give them inductive biases more similar to low-dimensional SSMs. Overall, current understanding of the community is that performant LCSMs are empirically indistinguishable from low-dimensional SSMs.

Thus I have significant doubts about the practical utility of the proposed algorithm for fundamental technical reasons. However, I understand that these empirical points are not comprehensive.

Even if we ignore these concerns about the applicability, which I am willing to accept on the grounds that the algorithmic contribution is solid and can be the potential basis for future work, I cannot increase my score without a significant improvement in presentation. In particular, it should discuss the expressivity of LCSMs with more accuracy/nuance and present the taxonomy of various LCSMs and associated tradeoffs much more carefully. These questions are central to the context, significance, and applicability of this work. Concretely, for example, it should benchmark the proposed algorithm for general LCSMs against the inference of recurrent LCSMs (and ideally other LCSMs e.g. based on dilated convolutions). Note that I am not saying that the paper needs to show that its algorithm is faster, but that it should show the crossover point and discuss the theoretical and practical tradeoffs accurately, to better position the proposed contributions within the broader space of closely related literature.

Thank you for the comments! We are happy to try and integrate them in our second revision where we will focus on a deeper contextual discussion. However, before doing so we were hoping to make sure we are on the same page about the current revision.

Currently, SSMs are barely mentioned in the it and these are the only places we could find that we thought can be misinterpreted or confusing and our (performed) fixes:

• At line 108 we could make it explicit that we mean parallel-scan-based training would be quadratic - we expected the word LTI SSM approach to convey this.
• At line 197-198, we will re-emphasize low-dimensional.

While we agreed and acknowledge the first rebuttal error (regarding stating S4 being the only one using FFT for training), we are not aware of any other places that you refer to through <incorrect statement perpetuated throughout the paper>, <often implied that "LCSM" does not include SSMs">, <many more technically incorrect statements>. We agree with the point that we could and should add more nuance and are working on that, but in the current form, the paper simply does not try to cover the differences and we do not see what technically incorrect statements you are referring to. Do you mind pointing us to what other places you believe we should fix in the current version? It will be helpful for our upcoming revision.

Regarding points:

• P1: We never meant to refer to LCSMs as "LCSMs that cannot be represented by a low dimensional SSM", but rather to emphasize that LCSMs contain strictly more than low-dimensional SSMs. Could you please pinpoint which part of the paper gave this impression of a dichotomy?
• P2: We were assuming to be in the context you mentioned (LTI and SISO). Granted this, it is a function of state dimension whether one chooses FFT or parallel-scan, both in theory and in practice. It is true that even the theory predicts the threshold to be at $\log(L)$ which once accounted for constant factors can be (and truly is) very small. However, we do not see how our statement is incorrect (again, assuming LTI and SISO as the reviewer suggested to assume) - we never defined what low means precisely because we only mean something below a certain threshold, which could indeed be very low. And while we agree to the point that for such SSM dimensions, one would prefer to use the convolutional view during training (and we did acknowledged this mistake in the rebuttal), as well as to the value of emphasizing it in our work, we are not aware of having made such claims in the paper (and thus do not see how this was perpetuated in the discussion).
• P3: we agree this could be useful to know and we will be trying to benchmark them, but since we need to have a fair comparison we need to integrate them in our current implementation. We will do our best to do so in due time.
• P4: the only point where we did not mention low-dimensional was lines 197-198 mentioned above. Do we agree or not that low-dimensional SSMs are not fully expressive in terms of the convolutions they can represent? While it is common to use underparameterized forms of LCSMs our definition uses the word often (line 160) to not constraint to that setting. In the introduction as well, at line 44, we mention they are just potentially underparameterized. So in the current state we do believe LCSMs could also be explicitly parameterized if one desires to do so (and, our method does not make any assumptions on how they are computed). This has (up to regularization) even been attempted by [1].
  • Furthermore, even if they were underparameterized (which again was not our intention to mean, nor do we think we did so, but we can emphasize this to avoid confusion), depending on the choice of parameterization, one could easily extend to filters that are not representable via low-dimensional SSMs. Any one filter has at least some (likely unrealistic) parameterizations that would cover it. That parameterization corresponds to a choice of bias.

However, we agree that there is no need to make an arbitrary distinction and that the method also applies to low-dimensional LTI SSMs. The reason for the distinction was simply because the low-dimensional SSMs already have a good (time and memory efficient) inference alternative. We appreciate the point that one could use our method for say average-dimensional SSMs where to determine the right dimensional threshold we would need the experiment you mentioned; we will touch on this.

[1] Daniel Y Fu et al. Simple hardware-efficient long convolutions for sequence modeling, International Conference on Machine Learning. PMLR, 2023.

Thank you for acknowledging the presentation and contribution of the paper, as well as for pointing out to the very useful references.

We have mentioned the multiresconv paper and how it has a linear time inference option in Appendix A.2. Granted that structure of a convolutional filter, this inference approach will always be preferable to ours efficiency-wise; we have however added more points about why having a broadly applicable fast LCSM inference method is valuable and can be more valuable in the future. This includes emphasizing the data-dependent setting which we consider a potentially relevant open problem.

On a similar note, we have also added a more nuanced discussion about how our method fits the bigger picture (the whole of Appendix A, with an in-depth analysis of the relationship between LCSMs and SSMs) and the tradeoff between computational efficiency and structural exploitation.