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

Authors: Costin-Andrei Oncescu, Sanket Purandare, Stratos Idreos, Sham M. Kakade
Primary area: generative models
Keywords: computational efficiency, inference, sequence generative models, long convolution sequence models, hyena

Abstract:
While transformers have been at the core of most recent advancements in sequence generative models, their computational cost remains quadratic in sequence length.
Several subquadratic architectures have been proposed to address this computational issue. Some of them, including long convolution sequence models (LCSMs), such as Hyena, address this issue at training time but remain quadratic during inference. We propose a method for speeding up LCSMs' exact inference to quasilinear time, identify the key properties that make this possible, and propose a general framework that exploits these. Our approach, inspired by previous work on relaxed polynomial interpolation, is based on a tiling which helps decrease memory movement and share computation. It has the added benefit of allowing for almost complete parallelization across layers of the position-mixing part of the architecture. Empirically, we provide a proof of concept implementation for Hyena, which gets up to $7.8\times$ end-to-end improvement over standard inference by improving $110\times$ within the position-mixing part.

Decision: accept (poster)
Reviewer scores: 8, 5, 8, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730592121185

thread_type:
official_review

<thread_object>
creator: reviewer_byxw

<content>
summary:
This paper develops a computation framework for efficient inference of long convolution sequence models (LCSMs). The proposed approach is based on relaxed polynomial interpolation, and can speed up LCSMs’ exact inference to the quasilinear time, as is shown by both complexity analysis and numerical verifications.

strengths:
1. The paper is well-written and easy to follow. 
2. The complexity analysis regarding both the computation and memory is detailed. 
3. The experiments are also provided to numerically verify the effectiveness of the proposed method.

weaknesses:
1. The main concern is about the novelty of this paper, particularly compared to the work by van der Hoeven and "dynamic FFT". It would be clearer if authors can provide more details about the differences or improvements of this work compared to former references. 
2. The applied setting seems somewhat limited, since it currently works only for Hyena-related architectures. Does the proposed method have the potential to inspire further extensions for general state-space models (SSMs; as with an equivalent convolutional filter) and Transformers (despite reasons stated in Sec. 3.4)? 
3. The numerical experiments part can be strengthened by e.g. adding tests for multiple configurations of hyper-parameters.

questions:
1. Please kindly provide more details to the questions raised in the "Weaknesses" section. 
2. How about the case when the non-mixers’ runtime is dominant, e.g. the large MLP module is common in practice (such as Transformers in applications). Do the improvements shown in Fig. 2(a) & 3(c) become marginal? 
3. For Fig. 3(c), why is the non-mixers’ runtime less for Eager (NP) & Lazy (NP)? In addition, it seems that there is no reduction of mixers’ runtime of Hybrid compared to (Flash) FFT?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730682147245

thread_type:
official_review

<thread_object>
creator: reviewer_pkhg

<content>
summary:
This paper introduces "Flash Inference", a novel framework for efficient inference in long convolution sequence models (LCSMs) employing a tiling strategy inspired by relaxed polynomial interpolation . The authors propose an algorithm that achieves quasilinear $O(L \log^2_2 L)$ time complexity for exact inference, improving upon the quadratic complexity of standard approaches without resorting to any approximation techniques. The work specifically demonstrates the method's effectiveness on Hyena architectures, achieving significant end-to-end speedup and even more significant improvement within the position-mixing component of the model. Beyond LCSMs, the authors identify key properties that enable such speedups and propose a general framework to guide future architecture design. The proposed approach further reduces memory movement and enables parallel computation across layers.

strengths:
1. Technical Innovation and Soundness:
- The paper presents a mathematically rigorous approach to improving inference efficiency, with clear proofs and careful analysis of complexity bounds
- The implementation details are thoroughly considered, including memory optimization and parallelization strategies
2. Practical Impact:
- The achieved speedups (1.6× end-to-end, 50× for position-mixing) are significant and well-documented
- The method is exact rather than approximate, maintaining model fidelity while improving performance
- The framework extends beyond just LCSMs, providing valuable insights for future architecture design
3. Experimental Validation:
- Comprehensive empirical evaluation across different hyper parameters parameters ($B$, $M$, $D$)
- Careful ablation studies of different $\tau$ implementations
- Clear breakdown of performance improvements and their sources

weaknesses:
1. Technical Presentation:
- The notation could be improved for clarity, particularly in handling subscripts that simultaneously indicate sequence position, feature dimension, and layer. For example, you could consider using superscripts to indicate features (channels) and/or paranthesis around the subscripts/superscripts for the layers. 
- the use of "$\mapsto$" notation on line 159 is imprecise; it should be $y\mapsto z$.
- Algorithm 1 should reference equation (3) for the definition of $\tau$ for better clarity.

2. Contextual Discussion:
- The discussion of approximate inference methods could be more nuanced. The statement about "defeating the purpose of using LCSMs instead of LTI SSMs" is overly strong, given that Hyena filters are intentionally underparameterized and some *implicit regularization* (convergence of the filters to ones that can be effectively represented by a truncated basis of complex exponential functions) is to be expected. Te reviewer suggests that the authors provide a more balanced discussion of the trade-offs between exact and approximate methods, considering the intentional underparameterization of Hyena filters and the potential benefits of implicit regularization.
- The introduction would benefit from a clearer distinction between prefilling and autoregressive generation phases in the complexity analysis. Elaborating on this could strengthen the clarity of how the method affects the performance of the system.
- The claim about "popular misconception" (line 109) regarding the realization theory of LTI systems could be better supported with appropriate citations. It is the reviewer belief that the this is not a misconception in the control theory literature as the problem of realization of LTI system has bee studied for more than 60 years  (e.g., Kalman and Ho 1960). Perhaps the authors could clarify on that.

3. Technical Specifics:
- The asymptotic complexity expressions should consistently specify the base of the logarithm ($log_2$)
- The parallelization discussion could benefit from more detailed analysis of the memory bandwidth implications
- The interesting extension to data-dependent filters, while covered in the appendix, could deserve more attention in the main text if supported by quality experiments (do we need input-dependent filters in LCSMs?)

questions:
1. How does the performance of Flash Inference compare to existing methods when dealing with very long sequences (e.g., $L>$32K) in non synthetic scenarios? Are there specific challenges (other than memory) or additional optimizations possible at these scales?
2. Could you elaborate on how the tiling strategy might be adapted for architectures with more complex dependencies between layers? This seems particularly relevant for hybrid architectures that combine convolution and attention mechanisms.
2. The main limitation of the proposed method is certainly the linear growth of cache size with sequence lenght. How does this affect the practical benefits of exact vs approximate inference methods? In the approximate method, the cost of distillation can be amortized offline for a subsequent constant memory autoregressive generation. Moreover, one can trade additional memory (number of states in the SSM) for additional precision in representing the filters. Using a pre-trained LCSM, could you verify the accuracy-memory trade-off of the approximate vs exact method for practical sequence lengths? 
3. Have you considered how this framework might be extended to handle structured sparsity patterns in the filters? This could potentially lead to additional efficiency gains while maintaining the theoretical guarantees.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731037056625

thread_type:
official_review

<thread_object>
creator: reviewer_qm9l

<content>
summary:
This paper proposes a method for speeding up Long Convolution Sequence Models (LCSM)'s exact inference to quasilinear time, identifies the key properties that make this possible, and proposes a general framework that exploits these. The proposed approach is inspired by relaxed polynomial interpolation and uses a tiling method that minimizes memory movement and enhances computation sharing, allowing near-complete parallelization across layers in the architecture’s position-mixing component. Through a proof-of-concept implementation on Hyena, we demonstrate up to a 1.6× improvement in end-to-end inference time and a 50× speedup in the position-mixing part.

strengths:
1. The paper addresses a fundamental issue in sequence models, particularly long convolution sequence models (LCSMs) like Hyena, where inference time scales quadratically with sequence length. The proposed framework reduces this to quasilinear time, achieving significant improvements by leveraging a tiling-based approach inspired by relaxed polynomial interpolation.

2. Besides speed, the paper focuses on reducing memory movement, a bottleneck in handling large models. It suggests methods for activation storage optimization and discusses adjustments for memory-constrained hardware, making the work relevant for deployment.

3. The paper presents thorough empirical results that demonstrate the proposed method's effectiveness. It reports up to 1.6x end-to-end improvement in speed for Hyena and a 50x speedup within the convolutional mixer component, which provides concrete evidence of practical impact.

weaknesses:
1. While the study is well-documented, the experiments focus primarily on synthetic setups and Hyena. Additional tests on a broader range of sequence models or with real-world tasks could further validate the framework’s generalizability.

2. The framework's efficiency relies on data-independent filters for optimal performance. Although data-dependent filters can be supported, doing so may lead to higher complexity or additional constraints. This could limit the application in cases where data-dependent filters are essential.

3. The framework assumes autoregressive causality, which could restrict its application to non-causal architectures. This point is mentioned briefly, but it would benefit from further exploration of how the framework could adapt to different model constraints.

questions:
1. The paper suggests applicability beyond LCSMs. Could the authors provide additional empirical results or theoretical discussion on adapting the method to different architectures, such as transformers, beyond high-level discussion?

2. While data-independent filters optimize the framework's efficiency, data-dependent filters could require modifications. Could the authors elaborate on practical adaptations to support data-dependent filters effectively, and whether the framework maintains similar efficiency in such cases?

3. The work compares primarily with the eager and lazy approaches. How does the proposed framework perform against other efficient models or architectures, such as those using efficient transformer variants?

4. The paper assumes causality in model design. Are there potential modifications to the framework that could apply it to non-causal architectures? This could expand the impact and applicability of the proposed method.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731100887261

thread_type:
official_review

<thread_object>
creator: reviewer_xhkl

<content>
summary:
This work introduces a method to speed-up the autoregressive inference of long convolution sequence models (LCSMs) to near linear time. The approach is based on findings from relaxed polynomial interpolation’s literature, which is gratefully adapted to convolutions.
The resulting algorithm results in important speeds-up both for the position-mixing components and end-do-end inference.

strengths:
- The method is novel and offers interesting improvements in the inference speed of LCSMs.

- The paper offers interesting perspectives that could be used for the design of more efficient (causal, input-dependent) LCSMs in the future.

weaknesses:
- The main weakness of the paper is that the presentation, design decisions and final implementation of the method remains quite abstract, even after reading the paper multiple times. Given that the paper presents an inference strategy, it should be feasible to have an stand-alone implementation (at least for one layer) incorporated in the Appendix of the paper. This would give clarity to the final, concrete version of the algorithm.

- Next, I feel that the presentation of the paper could be improved. For example, the function $\tau$ –which is crucial for the method– remains undefined through the whole body of the paper up to Sec. 4.2, where it is briefly and loosely defined. It is mentioned that 7 possible versions were tested, but only 4 were mention-worthy. Then, it is mentioned that some are used, but it is not specified in which ranges and under which parameters one is preferred over the other. Given that this is core to the method, this should definitely be improved.

questions:
### Additional questions and observations

- Contribution 1. This is true, however, for clarity, I would recommend that the authors mention Laughing Hyena before making this claim to put in context that that method is not exact.

- Line 233. There’s a typo here that changes the whole meaning of the sentence. Please fix.

 ### Conclusion

While I believe that the core contributions of this paper are very valuable, I am not really convinced by the current presentation of the paper. I would recommend the authors to improve readability and make specific the different design decisions that make the paper. 

Do note that the contribution of this paper is rather an algorithmic one. Yet, the current very abstract depiction of the method makes it unnecessarily difficult to implement and reproduce. I am, therefore, only able to provide this paper a “weak acceptance” at this time. However, that if the authors were to improve over the weaknesses outlined here, I would be more than happy to increase my score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731108780480

thread_type:
official_review

<thread_object>
creator: reviewer_irjb

<content>
summary:
This paper proposes more efficient algorithms for auto-aggressive inference of long convolutional sequence models (LCSMs). The aggregate running time on a sequence of length $L$ is reduced from $O(L^2)$ to $O(L \log^2 L)$, and the actual wall clock time of the implemented algorithm reflects substantial increases in efficiency.

strengths:
The paper's main strength is that it is technically sound and novel, and addresses the problem it aims to solve.
- The technical writing is clear and is helped by the inclusion of helpful graphics and rigorous algorithm boxes.
- Many considerations and variants of the core algorithm are proposed.
- An actual implementation is provided and all variants and baselines are benchmarked empirically.

There is a conception that long convolutions cannot be implemented efficiently in autoregressive inference settings, and so I do think that this paper presents an original algorithmic contribution.

weaknesses:
While the paper provides a technical contribution, the paper's main weakness is that of significance and direction with respect to the broader field; it aims to solve a problem that I believe does not need solving. Correspondingly, the papers writing (in terms of positioning and related works / baselines) could also use improvement.

- The paper's related work is sparse and I think it is important to present the lineage of these models more carefully. The original (depth separable) LCSMs were independently developed by two lines of work: the implicit convolution (CKConv and FlexConv) line of work, and the SSM line of work (LSSL, S4, DSS/S4D, H3, MEGA/Megalodon, and many more). Even though I understand why the paper deliberately puts emphasis on LCSMs that are not SSMs, because this is where its results are most applicable to, the positioning is at times misleading (e.g. in paragraph 2 of the introduction, where LCSMs are implicitly defined as non-SSM models, even though LTI SSMs are in fact the original models that popularized LCSMs).
- It is odd that the paper is heavily anchored around the Hyena architecture, even though it does not actually do anything Hyena-specific. The experiments don't use an actual trained model or look at empirical performance, only the speed of an architecture, for which any model with a LCSM convolution (e.g. H3, to which Hyena is equivalent except for the choice of convolution kernel) could equivalently be substituted into the writing without changing the algorithmic results.
- A distinction is made that low-dimensional LTI SSMs cannot represent general LCSMs, which is true and where the paper's potential benefits can come from. However, current understanding of LCSMs is that they need to be defined using certain priors (e.g. baking in exponential decay as in SGConv and Hyena) that essentially are similar to the priors imposed by low-dimensional SSMs. Empirically, it is known that there is little difference between these models (e.g. https://arxiv.org/abs/2312.00678v1 Fig 1, which claims that H3 in fact performs better than Hyena, where the only difference is the choice of convolution kernel parameterization). Thus a major weakness of this paper is that it is currently only applicable to models that have better alternatives empirically.

questions:
I think the paper needs to position with respect to other types of LCSMs with more nuance. For example, it could benchmark how the proposed algorithm for general LCSMs compares to the inference speed of recurrent LCSMs; e.g. for a given sequence length $L$, at what recurrent state size does the speed of a recurrent SSM cross over the speed of the proposed algorithm? This would at least provide some more useful context for the reader interested in the pure algorithmic aspects.

However, overall, in order for this paper to be valuable to the machine learning community, it should be applied to actual models, and current understanding of LCSMs is that general LCSMs essentially do not benefit over low-dimensional LTI SSMs. The contributions are interesting from a purely algorithmic perspective, and the paper does suggest potential extensions beyond time-invariant convolutions. However, it is not clear whether this has any chance of being extended into performant models - no actual downstream model is proposed. Thus, I think that a conference such as ICLR is not the most appropriate venue for such a submission without anchoring the algorithmic contributions to real models. In order to increase my score substantially, I think the paper needs to show a practical utility, for example either
- Showing that there exist classes of models that the algorithm can be applied to that are not dominated by others (e.g. showing that some non-SSM LCSM model is actually faster or more performant than an equivalent SSM)
- Showing that there exist data-dependent extensions that outperform data-independent convolutions in modeling performance
</content>
</thread_object>

</thread>
<thread>
creation date:
1731349740786

thread_type:
official_review

<thread_object>
creator: reviewer_zv8x

<content>
summary:
This paper introduces a novel linear-time inference algorithm for long convolution architectures like SGConv and Hyena. These models typically have a quadratic inference complexity with respect to the sequence length N, which can be prohibitive for large-scale applications.

The key insight behind the proposed approach is a clever partitioning and precomputation strategy for the contributions to the convolution of future outputs. This allows the inference complexity to be reduced to O(N log^2 N). The paper provides extensive experimental evidence demonstrating the acceleration of inference achieved by this new algorithm.

strengths:
* As far as  I know, the interpolation perspective presented in this paper is original and inspiring. The writing is exceptionally clear, and Figure 1 has been extremely helpful in understanding the proposed method.

* The algorithm introduced in the paper largely solves the long-standing problem of quadratic inference complexity for long convolution models like SGConv and Hyena. This has been a significant bottleneck for the practical deployment of these architectures (note that there are also other long conv architectures that do not suffer from this, please see below).

weaknesses:
This is a good paper and there is no much weakness to say about its methodology. However, I find the significance of the work depends on a line of work on long convolution architectures that the authors unfortunately have not discussed or compared.

Long convolution kernels can be contructed from smaller convolutions in a tree style hierarchical dilations such as those in WaveNet [1]. Recently, people have shown that these architectures, with nonlinearities removed and weight sharing, can be interpreted as having a wavelet based state and are competitive with SSM and long convs on sequence modeling benchmarks [2].

Crucially, these dilated convolution architectures also support linear-time inference, as detailed in [3], by maintaining a cache per layer. This allows for efficient inference without the need for techniques like FFT that are required for standard long convolutions.

Given these relevant prior results, I believe a comparison of the proposed approach to these dilated convolution models and their linear-time inference capabilities would greatly strengthen the paper. These alternative architectures offer a potentially simpler approach. As a reader of this paper, it would be nice to know when it is preferable to have a less structured long conv combined with the inference acceleration presented here.

References:

[1] Van Den Oord, Aaron, et al. "Wavenet: A generative model for raw audio." arXiv preprint arXiv:1609.03499 12 (2016).

[2] Shi, Jiaxin, Ke Alexander Wang, and Emily Fox. "Sequence modeling with multiresolution convolutional memory." International Conference on Machine Learning. PMLR, 2023.

[3] Paine, Tom Le, et al. "Fast wavenet generation algorithm." arXiv preprint arXiv:1611.09482 (2016).

questions:
Please see my question above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732350281959

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
First Revision (including vastly improved experiments)

comment:
We provide an intermediate revision of our paper. 

The primary experimental improvement we introduce is the use of CUDA Graphs to address CPU overhead in non-mixer components during inference. In both synthetic settings and the Hyena model, these components primarily deal with only current token being generated, where GPU compute time is often shorter than kernel dispatch time. While this overhead is typically hidden by the mixer’s GPU compute time, Flash Inference reduces mixer times substantially (up to $120\times$), exposing the CPU overhead and decreasing end-to-end speed-ups. Lazy execution also exposes this overhead due to its faster performance compared to Eager.
CUDA Graphs resolve this by recording all kernel dispatches for a single token generation into a graph and replaying it for subsequent tokens, dispatching all non-mixer kernels simultaneously. This significantly reduces CPU overhead, making it constant for Flash Inference and its baselines. As a result, Flash Inference achieves up to $120\times$ speed-ups in the mixer and up to $7.8\times$ in end-to-end Hyena runs. A more detailed discussion can be found in Appendix D2 (CPU Overhead: The silent performance killer).

Regarding the points of the reviewers:
- we are currently working on adding a more comprehensive discussion on related techniques and an extended literature review.
- we will also do our best to perform experiments comparing our method to inference for LTI SSMs of various dimensions.

We tried to address every other comment (and referenced this in the corresponding individual rebuttals) except for the ones related to these 2 points, as follows:
- R2W1: Modified lines 108 and 198-199 to emphasize that what we consider a (strict) subset of LCSMs is not SSMs but rather low-dimensional SSMs. We’ll do more about this.
- R2W2: We dropped the references to Hyena at lines 17, 43, 213 and 530 - now it is only referenced for experiments purposes and as an illustrative example at line 165,
- R3W1: need an algorithm
- R3W2: we added an explicit definition of $\tau$ just before Lemma 1 rather than within it. We also mentioned the other 3 implementations in the Appendix (under “More Implementation Details”).
- R3Q1: We have added a mention to Contribution 1 that Laughing Hyena is an approximation rather than exact inference.
- R3Q2: We changed the confusing sentence (former line 233)
- R4W2: added a clear mention to the extra factor of 2 for data-dependent filters at line 214.
- R4[W3/Q4], R5Q2: we added a discussion on causality in Appendix “Framework Adaptations”
- R5W3.2: we added a longer discussion on the connection between across-layer parallelization and memory bandwidth in Appendix “Discussion on Memory”
- R5Q2: we added a discussion on hybrid architectures in Appendix “Framework Adaptations”
- R6W3/Q3: we added comprehensive and easily inspectable data on all our runs (21 experiments) in Appendix Further Experimental Results
</content>
</thread_object>

</thread>
<thread>
creation date:
1732444587669

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Second Revision (added an extensive contextual discussion)

comment:
We have fully updated our background review: updated lines 103-112, 197-199 as well as wrote appendix A on an in-depth contextual discussion. Particularly:
- We changed the perspective to treat LCSMs and SSMs equally and only when relevant for computational reasons distinguish low-dimensional SSMs (as suggested by Reviewer 2)
  - Thus, also emphasized our applicability to SSMs as well. Here, theory predicts that our method will be faster for state dimensions of $\Omega(\log^2(L))$ - however, there is a memory cost associated with it. We are doing our best to get experiments confirming this done in due time.
- Besides SSMs, also commented on how multiresolution convolutions represent another structured class of convolutions where better inference methods exist, as pointed by Reviewer 1.
- Added more comments on the distinction between prefilling and autoregressive generation.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a method for speeding up long convolution sequence models exact inference to quasilinear time. The paper provides extensive experimental evidence demonstrating the inference benefits of this new algorithm. The reviews for the paper were mostly positive with reviewers highlighting the significance of the work and novelty. I recommend acceptance.