Title: Laughing Hyena Distillery: Extracting Compact Recurrences From Convolutions
Venue: NeurIPS 2023
Accepted: yes
Keywords: Long convolutions

recurrence

attention

language models

signal processing

throughput

auto-regressive generation
Abstract: Recent advances in attention-free sequence models rely on convolutions as alternatives to the attention operator at the core of Transformers. In particular, long convolution sequence models have achieved state-of-the-art performance in many domains, but incur a significant cost during auto-regressive inference workloads -- naively requiring a full pass (or caching of activations) over the input sequence for each generated token -- similarly to attention-based models. In this paper, we seek to enable $\mathcal O(1)$ compute and memory cost per token in any pre-trained long convolution architecture to reduce memory footprint and increase throughput during generation. Concretely, our methods consist in extracting low-dimensional linear state-space models from each convolution layer, building upon rational interpolation and model-order reduction techniques. We further introduce architectural improvements to convolution-based layers such as Hyena: by weight-tying the filters across channels into heads, we achieve higher pre-training quality and reduce the number of filters to be distilled. The resulting model achieves 10x higher throughput than Transformers and 1.5x higher than Hyena at 1.3B parameters, without any loss in quality after distillation.

LAUGHING HYENA DISTILLERY Extracting Compact Recurrences From Convolutions

Abstract

Recent advances in attention-free sequence models rely on convolutions as alternatives to the attention operator at the core of Transformers. In particular, long convolution sequence models have achieved state-of-the-art performance in many domains, but incur a significant cost during auto-regressive inference workloads - naively requiring a full pass (or caching of activations) over the input sequence for each generated token - similarly to attention-based models. In this paper, we seek to enable O (1) compute and memory cost per token in any pre-trained long convolution architecture to reduce memory footprint and increase throughput during generation. Concretely, our methods consist in extracting low-dimensional linear state-space models from each convolution layer, building upon rational interpolation and model-order reduction techniques. We further introduce architectural improvements to convolution-based layers such as Hyena : by weight-tying the filters across channels into heads , we achieve higher pretraining quality and reduce the number of filters to be distilled. The resulting model achieves 10 × higher throughput than Transformers and 1 . 5 × higher than Hyena at 1 . 3 B parameters, without any loss in quality after distillation.

1 Introduction

Attention-free approaches such as long convolution sequence models (LCSMs), e.g., H3 [1], Hyena [2], have shown promise in matching Transformer [3, 4] performance across a wide range of tasks, with sub-quadratic complexity with respect to sequence length. Despite the improved efficiency during training on long sequences, unless the convolution filters are either short or admit a low -dimensional state-state-space realization, LCSMs still need to process the entire growing sequence at every step of auto-regressive generation, similarly to Transformers.

In this work, we seek to refine LCSMs in both efficiency and quality . First, we study the inference stage, and propose methods to enable a recurrent mode for auto-regressive generation. Recurrent modes prescribe the existence of a state encoding the past information of the process in a fixed-dimension memory, enabling constant per-step time and constant-memory in generation. Then, we draw upon an analysis of pre-trained models to develop architectural enhancements for the Hyena block, simultaneously improving model quality and efficiency of the distillation procedure.

Distilling fast recurrences We introduce LaughingHyena , the first distillation approach for LCSMs that enables recurrent inference without impacting downstream quality. LaughingHyena seeks compact recurrences in the form of state-space models (SSMs) [5, 6] as the solution of a nonlinear interpolation problem involving the convolution filters of a pre-trained model. Since the total memory cost of SSMs grows linearly in the state dimension d , our distillation procedure enables high throughput by enabling processing of large batches during generation. We identify and address three core challenges related to distillation, including the identification of:

Target state dimension: we identify candidate state dimensions of our distilled SSMs by analyzing the spectrum of the Hankel operator associated with each convolution [7].

Parametrization: we address issues with naive parametrizations by introducing a factorized modal form, inspired by barycentric [8] and Prony-like [9] methods .

Approximation metric: to ensure compatibility with any downstream task, we choose discrepancy metrics on the convolution filter, rather than model outputs.

In auto-regressive workloads, LaughingHyena -distilled models with state dimension d can generate K tokens in O ( dK ) time and with constant O ( d ) memory - improving over the O ( K 2 ) time and O ( K ) memory usage of kv-cached Transformers and naively executed long convolutions. At model sizes above one billion parameters, LaughingHyena achieves 10 × higher peak throughput over comparable Transformers (Figure 1.1), and can process larger batch sizes. Constant memory generation enables larger K for a given a memory constraint e.g., generating 512 tokens with LaughingHyena requires 3 × less memory than with a Transformer. At smaller batch sizes, latency of LaughingHyena is also competitive with Transformers, reaching ≥ 2 × speedups at longer prompt lengths.

Improving pre-training quality We leverage our analysis of the distillation process to open up new avenues of improvement for LCSM architectures. Indeed, the high compression rates achievable through LaughingHyena hint at subutilization of the convolution. We revisit the multi-headed design of H3 [1]; tying weights across channels pushes long convolution filters towards larger effective dimension, and as an additional advantage reduces the runtime of posttraining distillation and inference memory footprint. Further, multi-head Hyena models improve on pre-training perplexity over regular Hyena and GPT [10] architectures on the language dataset THE PILE [11].

2 Preliminaries and Related Work

100

200

300

400

500

0

.

2

0

.

4

0

.

6

0

.

8

1

·

10

4

Generation Throughput

[

tok

/s

]

Laughing Hyena 1

.

3B

Hyena 1

.

3B

.

H3 1

.

3B

Transformer 1

.

3B

We discuss convolutions, state spaces and auto-regressive generation workloads for sequence models.

Convolutions Let ∗ denote the convolution operator. It is defined as the dual operation to pointwise multiplication under Fourier transform. In signal processing and deep learning alike, one often encounters the causal linear convolution of a filter h (which may extend indefinitely) with an input u of length L :

Generally, u t ∈ R D where D is the width of the signal - or in deep learning parlance - the number of channels . Without loss of generality, we specialize our analysis to single input single output layers, i.e. with D = 1 . For the input-output relations of type (2.1), we use the terms convolution layer and linear system interchangeably. Similarly, the function t ↦→ h t is referred to as both the filter and the impulse response of a linear system. Existing convolution sequence models can be classified in terms of the parametrization used for their filters. The class of implicit convolutions represent the filter as a parametric function γ θ : t ↦→ h t .

State-space realization One option is to select γ θ as the impulse response function of a discrete linear time-invariant system,

with state x t ∈ R d , input u t ∈ R , and output y t ∈ R . The matrices A ∈ R d × d , B ∈ R d × 1 , C ∈ R 1 × d , and h 0 ∈ R are the learnable parameters of the model while the initial state x 0 is usually set to zero such that u ↦→ y is a pure convolution. While linear systems (2.2) are the staple of signal processes and control theory, their use as implicit parametrization of convolution filters in deep neural networks have only recently emerged [12, 6]. Other parametrizations [13, 14, 2] select γ θ ( t ) as different flavors

of implicit representation neural networks [15, 16]. The latter are generally more powerful in terms of the class of filters they can represent and flexibility during training, at the cost of losing a fixed state dimension. q

2.1 Long Convolution Sequence Models

The H -family of convolution sequence models H3 [1] and Hyena [2] - relies on a combination of long convolutions and datacontrolled gating to replace attention with sub-quadratic scaling in sequence length 1 . We use the deep learning convention of naming different projections as query q , key k and value v . Let M q and

u

T

(

v

)

T

(

k

)

T

(

q

)

M

k

T

h

M

q

k

v

y

M k be the L -byL diagonal matrices whose respective main diagonal entries are the respective entries of lengthL sequences q and k . A H -block realizes a surrogate attention matrix with a data-controlled, parameterized decomposition in three terms:

where T h ∈ R L × L is the Toeplitz matrix constructed from the learnable long convolution filter h , i.e., T h =( h i -j ) L -1 i,j =0 . The qkv -projections are themselves the output of a convolution between the input sequence and three distinct short filters. The degrees of freedom in H -block design are the three short filters 2 and the long filter h . The long filter can be parameterized using an implicit neural representation [2], state-space model [1], or explicit values [17]. The threefold decomposition of the attention operator, allows evaluation of (2.3) in just ˜ O ( L ) := O ( L log 2 L ) time (two convolutions 3 and two element-wise products), y t = q t ( h ∗ kv ) t . The overall operator acts on an input u by constructing a thirdorder multi-variate polynomial of u whose coefficients are controlled (nonlinearly) by parameters of the block.

2.2 Auto-Regressive Generation

A typical workload for sequence models is auto-regressive generation. Given a lengthT prompt u ∈ R T , the model is tasked with producing the following K additional outputs - one at a time - for a resulting output sequence y of length L = T + K .

Convolution sequence models After processing the initial prompt in ˜ O ( T ) time and obtaining a lengthT output u ↦→ y 0 , . . . , y T -1 , a generic convolution layer can cache the output sequence and generate any additional outputs using (2.1) auto-regressively, i.e. y t +1 = ∑ t j =0 h t -j y j for t = T -1 , . . . , T + K -1 . It is important to note that auto-regressive generation with generic long convolutions is expensive. It comes with a quadratic cost in the number K of tokens to be generated and require storing a cache of length up to L .

Lemma 2.1. Generating K tokens with a long convolution layer (2.1) from a lengthT prompt has time complexity O ( T log 2 T + TK + K 2 ) and requires O ( L ) memory.

State-space models When the linear system admits a state space realization (2.2), i.e. it is able to switch between convolution and recurrent mode, the cost of auto-regressive generation can be dramatically reduced. The memory footprint is O ( d ) : all we need to cache is the state x t , a d -dimensional vector. With some further machinery that we develop in next section, we can retain ˜ O ( T ) time and O ( T ) memory to process the prompt 4 and initialize the state x T -1 . Each additional generation step only requires O ( d ) time.

Note that long filters h truncated to length d (i.e. h t =0 for t > d -1 ) can also be interpreted as d -dimensional SSMs (see Appendix A.7) where the state (a cache) coincides with the last d inputs.

Lemma 2.2. Generating K tokens with a state-space model (2.2) from a lengthT prompt has time complexity O ( T log 2 T + dK ) and requires O ( T + d ) memory.

Transformers Self-attention is certainly less efficient than long convolutions in processing the prompt, coming with a hefty O ( T 2 ) time complexity. However, Transformers can achieve a similar

efficiency in auto-regressive generation by caching the sequences of past keys { k t } and values { v t } . Specifically, from t = T -1 onward, the new projections ( q t +1 , k t +1 , v t +1 ) are evaluated from the current output y t , and the new output y t +1 can be computed in linear time with two reductions

Lemma2.3. Generating K tokens with self-attention from a lengthT prompt has time complexity O ( T 2 + TK + K 2 ) and requires O ( L ) memory.

3 The Laughing Hyena Distillery

In this section, we introduce our distillation method. We discuss choosing an approximation objective, a parametrization for the approximant and setting a target state dimension. Given any pre-trained LCSM, the objective of the distillation procedure is to convert each pre-trained convolution filter into a distinct state-space model (2.2). This should be achieved with the smallest state dimension d which preserves, up to a certain tolerance, the input-output characteristics of the convolution layer. Formally, given a filter h the distillation problem is defined as follows.

Given the sequence h 1 , . . . , h L , find a state-space model (2.2) of dimension d ≪ L , whose input-output behavior approximates the one of the convolution with h over the largest class of input sequences.

The choice of approximation metrics and assumptions on the input sequences yield different distillation objectives . A distillation algorithm constitutes a systematic procedure for optimally choosing the systems matrices with respect to a particular objective. In instances where the original filter h is itself the impulse response of a finite-dimensional state-space model, e.g., when attempting distillation of H3 or S4 [6] filters, the term distillation becomes analogous to model-order reduction . Hence, in such cases, the distillation algorithm should yield a state-space representation of a lower order state-dimension.

There exist several algebraic solutions to the model reduction problem [18, 19, 20], typically seeking low-rank structures of the state space by inspecting some invariant of the system, e.g. the Gramians in balanced truncation [19, Ch. 7]. The lower-order system is then obtained as a projection of the system dynamics onto the found subspace where the system retains desired characteristics, e.g., input-output behavior, stability, etc.

Truncated filters In theory, implicitly parameterized convolution filters can represent arbitrarily long signals. In practice, these filters are trained on a fixed maximum length L . At inference time the model can then be evaluated for sequences longer than L . During distillation it is nonetheless reasonable to treat the pre-trained filters as potentially very long (even beyond L ) but finite impulse response functions [21, 22, 23, 24]. We show how this choice is supported by empirical evidence displaying how pre-trained filters typically decay to zero in finite time (see Appendix D).

Transfer function representation Analternative description of the system (2.2) is its transfer function H , defined as the z -transform of the impulse response H ( z )= ∑ ∞ t =0 h t z -t for all z ∈ C where the sum converges. The transfer function is a proper rational function of z

In the z -domain, the transfer function defines the input-output map as Y ( z ) = H ( z ) U ( z ) . Here, H ( z ) is defined outside the C -plane circle of radius ρ ( A ) , D ρ ( A ) := { z ∈ C : | z | > ρ ( A ) } where ρ ( A ) is the spectral radius of A , i.e. the amplitude of its largest eigenvalue. We can recover all characteristics of a given system equivalently from either its transfer function or state-space representations (see Appendix A.3 for further details and derivations). Notably, the transfer function is an invariant of the system: if we apply a change of variables to the state, the transfer function remains unchanged (Lemma A.3). This alone should discourage attempts at modeling filters by learning dense state-space matrices A , B , C as such: there are infinitely many equivalent statespace realizations that map to the same system. Starting from coefficients ( a i ) and ( b i ) of the rational transfer function (3.1), we can compute the impulse response in ˜ O ( L ) time (Lemma A.6). Moreover, we can map back the transfer function to a special state-space realization - the companion canonical form - whose recurrence has time complexity O ( d ) (Lemma A.7), compared to

Sequence

Dataset

Pretrain

Conv

.

Model

h

Hankel

SVD

candidate

state dim

d

Distill

State

-

Space

A

,

B

,

C

, h

0

Compose

Recurrence

Compute

Transfer Function

Deployment

-

Downstream eval

.

-

Fine tuning

Get initial value

h

0

the O ( d 2 ) of dense state-space matrices. From Lemmas A.3 and A.7 we can also prove that any stable state-space model can be converted by canonicalization into its companion form, and thus can be equipped with an efficient recurrence (Thm. A.8).

The distillation problem presents several challenges:

Defining the distillation objective. A primary decision involves selecting a distillation objective. We are primarily interested in metrics of pure discrepancy between each filter of a pre-trained deep model and its approximator, rather than the expected input-output loss over a distribution of inputs.

Choosing a state-space parametrization. It is crucial to determine a suitable parametrization of the distilled state-space realization. Once this is decided, the task is to identify the parameters that minimize the distillation desiderata, which can involve challenging optimization problems in itself.

Selecting the target state dimension. Lastly, a challenge is to estimate the degree to which the model's order can be reduced. In other words, we must select the target state dimension of the distillation process to identify the right trade-off between efficiency and accuracy.

In the following, we address each of these challenges, and provide a comprehensive approach (summarized in Figure 3.1) to distill recurrences from convolution-based architectures.

3.1 Data-Free Distillation Objectives

We focus on distillation objectives that are independent of the training data and the overall architecture of the neural network under consideration. The distillation loss should be chosen as a pure measure of discrepancy between each convolution filter h t of the model and their finite-dimensional approximations ˆ h t = CA t -1 B . This approach ensures that we do not require a full sequential inference pass over the pre-trained model at each step of distillation procedure and the distilled model can be more broadly applied to downstream tasks. This choice is supported by Young's convolution inequality [25, 26], which indicates that the output approximation error has a bound ∥ y -ˆ y ∥ r ≤ ∥ h -ˆ h ∥ q ∥ u ∥ p for properly chosen norms 5 . For maximum numerical stability and freedom of parametrization for the approximants, we favor modern unconstrained gradient-based approaches to then solve the resulting distillation program 6 . We design distillation algorithms which either match filters in time domain minimizing the ℓ 2 error ( ∥ h ∥ 2 := [ ∑ t ∈ Z | h t | 2 ] 1 / 2 ) or match their transfer functions optimally with respect to the H 2 norm ( ∥ H ∥ 2 := [(1 / 2 π ) ∫ π -π | H ( e iω ) | 2 d ω 1 / 2 ) 7 . As the distillation is carried out via gradient methods, ℓ 2 is a natural candidate. H 2 error minimization can instead be used to uniformy bound the worst-case discrepancy as ∥ h -ˆ h ∥ ∞ ≤ ∥ H -ˆ H ∥ 2 (see Appendix A.2 for further details).

3.2 Making Hyena Laugh with Modal Interpolation

Our degrees of freedoms to solve the distillation problem are the matrices A , B , and C of the state-space realization, which determine the filter for all t > 0 . In distilled SSMs, the passthrough (residual) term cannot be freely assigned: it is simply h 0 , the value of the original filter at zero. Alternatively, given its

appealing invariance properties, we can parametrize a proper rational function ˆ H ( z ) (3.1) and fit it to the (truncated) transfer function 8 of the original filter H L ( z ):= ∑ L t =0 h t z -t (see Appendix B.2).

Modal canonical form Optimizing the full transfer function can be numerically challenging for several reasons e.g., ensuring stability 9 , and ill-posedness for high-order polynomials. A natural solution, inspired by barycentric approaches to rational function approximation [29, 8], is to assume d distinct roots λ n in the denominator's polynomial, λ n ∈ roots ( poly ( a )) .

Proposition 3.1 ([5]) . If poly ( a ) has distinct roots { λ n ∈ C } , then the transfer function of the system can be factorized as ˆ H ( z )= ∑ d n =1 R n / ( z -λ n ) , ∀ z ∈ D ρ ( A ) where { R n ∈ C } is the residue associated with the pole λ n .

Computing the inverse transform of the expanded transfer function via, e.g., the Cauchy residue theorem [30], shows that the resulting impulse response ˆ h corresponds to a truncated basis of exponentially decaying complex sinusoids

In practice, this corresponds to the impulse response of state-space model with diagonal matrix A = diag( λ 1 , . . . , λ d ) and such that B i C i = R i for all i = 1 , . . . , d . The distillation problem can be then defined in terms of the L -point nonlinear least squares interpolation error (squared ℓ 2 ) between h 1 , . . . , h L and (3.2) evaluated for t =1 , . . . , L : min { λ n ,R n } ∥ ˆ h -h ∥ 2 2 . Note that in case of the target filter h being real-valued, the objective can be replaced by ∥ R [ ˆ h ] -h ∥ 2 2 .

Modal Interpolation

0

20

40

60

80

100

120

-

0

.

5

0

0

.

5

1

h

t

R

n

λ

t

-

1

n

h

t

ˆ

h

t

(discrete) Fourier transform before computing the objective. Efficient evaluation of (3.2) is crucial for distillation. In particular we show the following:

Although we find solutions of the distillation (interpolation) problem via modern gradient-based optimization techniques, it is worth mentioning that Prony showed how the nonlinear least square solution can be computed solving two linear problems [9]. However, similar to Padé's method for rational approximation [31], these techniques can be numerically unstable. We opt for a parametrization similar to [32, 33] where each eigenvalue is parameterized in polar form λ n := A n e iθ n and the residues in cartesian form 10 . Note that, with this parametrization we have R [ ˆ h t ] = ∑ n A t -1 n [ R ( R n ) cos( θ n ( t -1)) -I ( R n ) sin( θ n ( t -1))] . We can also solve the distillation problem in the H 2 sense by evaluating ˆ h t and h t at t = 0 , . . . , L -1 and taking their respective

Lemma 3.1. Evaluation of ( ˆ h t ) L -1 t =0 (3.2) can be done in O ( dL ) time from its modal form and in ˜ O ( L ) time from its proper rational form.

3.3 Minimal Distillation Orders

Distilling into lower-dimensional systems is always desirable as they require fewer parameters to be optimized and they yield recurrences that are (linearly) more efficient in terms of time and memory complexity in post-distillation auto-regressive inference workloads. The dimension of the smallest possible state-space model with impulse response exactly { h t } t ∈ N is the so-called McMillan degree [34]:

Theorem 3.1 (Ho-Kalman [35, Theorem 2, Corollary]) . Let S be the (infinite) Hankel matrix constructed with h , i.e. S := ( h i + j ) ∞ i,j =1 . Then, d ⋆ = rank( S ) .

A lower bound for d ⋆ can be estimated from a truncated filter of length L by constructing the L × L principal sub-matrix S L and using the fact that rank( S ) ≥ rank( S L ) . Inspecting how fast the Hankel singular values ( σ n ) L n =1 decay in pre-trained convolution models can be predictive of the approximation quality at a fixed dimension. As a rule of thumb, d needs to be sufficiently large for σ d +1 to be sufficiently small 11 . Specifically, we can prove that the last singular value σ d determines the upper bound of distillation quality with a SSM of dimension d , in terms of the Hankel norm [19]. This is a direct consequence of Adamjan-Arov-Krein theorem [7] and can be informally stated as follows.

Theorem 3.2 (Informal) . Let h be a lengthL filter , ˆ h a distilled filter of order d < L and let S L , ˆ S L be the respective Hankel matrices. Then inf ˆ S L ∥ S L -ˆ S L ∥ 2 = σ d .

3.4 Deploying the Recurrence

Once all the filters of a pre-trained model have been distilled with the proposed modal interpolation technique described above, the model unlocks a recurrent mode which allocates a state x t ∈ C d for each filter and enables fast auto-regressive inference. Deployment of distilled model involves two critical steps: the pre-filling and the recurrent update rule itself.

Fast pre-filling During auto-regressive generation, when a lengthT prompt is fed to the model, we need to compute the state x T to start generating new tokens. Using the recurrence, the time complexity of initializing x T would be O ( dT ) with a O ( d ) memory footprint. One can alternatively distribute the computation on d processors with a parallel scan operation [37, 38] to reach a parallel time complexity O ( d log 2 T ) while incurring in an increased memory requirement of O ( dT ) 12 . A third option is to use a single FFT convolution to obtain x T in ˜ O ( T ) time and O ( T ) memory.

Proposition 3.2. x T = ( ν T , . . . , ν T -d ) where ν t = ( g ∗ u ) t and g is the filter whose transfer function is 1 / den ( ˆ H )( z ) and can be evaluated in ˜ O ( T ) .

Note that, the fast pre-filling algorithm established by this result requires evaluating the denominator polynomial of ˆ H from its roots before deployment. This is equivalent to converting the transfer function from its factorized representation to its rational form (3.1).

Recurrent step The update rule is diagonal, thus efficiently evaluated in O ( d ) time and memory:

Proposition 3.3. The filter (3.2) has a state space matrices A = diag ( λ 1 , . . . , λ d ) ∈ C d × d , B = (1 , . . . , 1) ⊤ ∈ C d × 1 , C = ( R 1 , . . . , R d ) ∈ C 1 × d whose step can be evaluated in O ( d ) time and memory.

As we generally want the output y t to be real-valued, we can simply update the complex state x t +1 = A x t + B u t and then take the real part of the output, y t = R [ C x t ] + h 0 u t .

4 Multi-head Long Convolutions

We can leverage the Hankel spectrum analysis discussed in Section 3.3 to study the dynamics of the effective dimensionality of each convolution filter during LCSMs pre-training. We find that, at initialization, filters correspond to high-dimensional SSMs, and gradually converge to lower-dimensional representations during training. See Appendix E.2 for examples on Hyena and H3 models.

This observation leads to the question: is it advantageous to perform independent long convolutions on each channel, or can we reduce the total number of filters without loss in quality? To answer this, we adapt the multi-head layer design proposed by H3 [1] to Hyena [2]:

Given the projections q, k, v ∈ R L × D , we split them into M chunks of size N = D/M , q m , k m , v m ∈ R L × N .

Each chunk is processed by a modified Hyena operator: first, we perform the outer product of k m and v m along the spatial dimension, z m := k m ⊗ v m ∈ R L × N × N , apply a long convolution with filter h m to all N × N elements independently, then compute y t m =( h m ∗ z m ) t q t m , y m ∈ R L × N as shown in Figure 4.

Finally, we compose y 1 , . . . , y m into a single output y ∈ R L × D via concatenation.

An instance of a MultiHyena is equipped with M < D distinct long convolution filters, which leads to ( a ) faster distillation, with less filters to approximate, ( b ) lower memory footprint, via a total reduction of the states to cache during generation and ( c ) faster filter generation, by tying the weights of filter parameters. We note that tying weights of key-value projections has also been shown to be an effective technique to reduce memory cost in Transformers [39, 40].

v

m

k

m

q

m

outer

T

m

h

dot

Theorem 4.1. The MultiHyena layer, with O (log s ) heads and model size O ( √ s log s ) can solve the associative recall problem, where s denotes the vocabulary size.

Crucially, the multi-head structure of MultiHyena enables us to prove favorable scaling in the associative recall synthetic task, which was shown in [2] to be predictive of performance at scale. In associative recall, the model is given a sequence of key-value pairs and a query, and is tasked with matching the query to a key in the sequence by returning its associated value. The difficulty of the task grows with the vocabulary size s : larger vocabularies necessitate wider models.

In Appendix E.1, we empirically verify improved scaling in vocabulary size with multiple heads.

5 Experiments

Pretraining: We pretrain a suite of MultiHyena language models on The Pile [11], investigating scaling of perplexity with different amounts of total tokens (5, 10, 15 billion), as well as larger training runs for 300 billion tokens. MultiHyena outperforms Transformers and Hyena .

Distillation analysis: We investigate the relation between optimal distillation orders, Hankel spectrum, and errors on the logits of distilled models.

Post-distillation downstreams: Weevaluate the downstream impact of distilling long convolutional language models, reporting HELM [41] and LM-Eval-Harness [42] results.

Benchmarking: We benchmark latency, throughput and memory along the different axes of batch size, sequence length, number of generated tokens. We include base models, distilled models and equivalent Transformers.

5.1 Pre-training

To validate the multi-head formulation, we train 150 and 350 million parameter MultiHyena models on The Pile [11] using 8 heads and otherwise the same architecture as equivalent Hyena models, following the setup of [2]. Via the multi-head structure introduced in 4, MultiHyena outperforms both Hyena and Transformers, including on data scaling runs with increasing numbers of tokens and full 300 B tokens runs (Table 5.1).

5.2 Distillation Analysis

Next, we verify whether Hankel singular values are predictive of downstream errors, and whether large models can be distilled without loss in quality. We apply LaughingHyena distillation to pre-trained MultiHyena , Hyena and H3 of different sizes. Concretely, for each layer and channel of a model, we parametrize the poles { λ n } of the modal canonical forms (Section 3.2) at different orders d , and solve for each ℓ 2 approximation problem.

Approximation errors and spectrum We investigate the magnitude of approximation errors introduced by LaughingHyena distillation. Given a pretrained MultiHyena model, we compute the errors between original and distilled filters at each layer, averaged across channels. We repeat this process for different distillation orders (state dimension of the model form of Section 3.2). Figure 5.2 visualizes

1

10000

20000

30000

40000

50000

-

1

0

1

2

error

Relative

/lscript

1

Error on Logits

d

= 8

1

200

400

600

800

10

-

1

10

0

Error Inset

d

= 8

1

10000

20000

30000

40000

50000

-

1

0

1

2

Sorted Logits

error

d

= 16

1

200

400

600

800

10

-

3

10

-

2

10

-

1

Sorted Logits

d

= 16

minimum, maximum and average errors, per-layer errors and the distribution of the singular values of the Hankel operator associated to each filter. We observe distillation orders ( > 16 ) that yield smalls errors to be predicted by the distribution of singular values. Thus, analysis of the Hankel operator's spectrum is verified to be an effective approach to direct estimation of the optimal distillation order. We also note that the optimal order changes across layers, offering options for further optimization.

Output errors Next, we compute relative ℓ 1 error between output logits of pre-trained and distilled models to ensure LaughingHyena can be used in generation workloads. The optimal minimal distillation order estimated via Hankel operators ( 16 ) is sufficient to keep the output distribution over the vocabulary ( > 50 k entries) close to the pre-trained model, as shown in Figure 5.2. Inspecting the error profile over logits sorted by magnitude reveals our approach to be robust to different sampling strategies for generation, including greedy decoding, topk , topp [43]. Indeed, the relative errors are < 10 -2 up to and including the 99 . 99% percentile of the distribution, meaning e.g., a topp sampling strategy with large p can be used on a distilled model without drift in outputs (mis-classified tokens). We note that the relative errors are maximum on small-norm logits, which are not required by most sampling strategies. In Appendix D.2, we provide a similar distillation error analysis for Hyena and H3 models. We find that Hyena and can be distilled with less than 32 orders and H3 with less than 8 .

5.3 Downstream Evaluation

We check how distillation affects downstream performance on language benchmarks. We apply distillation of order 8 , 16 and 32 to our THE PILE-pretrained MultiHyena language model and benchmark (Table 5.3) its performance on a suite of canonical (zero shot) tasks from LM-Eval-Harness [42] and HELM [41]. The results are consistent with our error analysis: distillation orders equal or greater to 16 introduce little-to-no quality degradation.

2 1

2 3

2 5

2 7

0

0

.

1

0

.

2

error

/lscript

1

Error

2 1

2 3

2 5

2 7

0

1

2

3

error

/lscript

2

Error

2 1

2 3

2 5

2 7

0

10

20

error

/lscript

∞

Error

2 0

2 2

2 4

2 6

0

2

4

6

·

10 2

singular value

Hankel Singular Values

2

3

4

5

6

8

10

12

16

20

32

64

128

1

3

5

7

9

11

13

15

17

order

depth

0

.

1

0

.

2

2

3

4

5

6

8

10

12

16

20

32

64

128

1

3

5

7

9

11

13

15

17

order

depth

0

1

2

3

2

3

4

5

6

8

10

12

16

20

32

64

128

1

3

5

7

9

11

13

15

17

order

depth

5

10

15

20

4

16

32

64

1

3

5

7

9

11

13

15

17

singular value index

depth

0

25

>

50

5.4 Benchmarking

We measure throughput, latency and memory usage of LaughingHyena for auto-regressive generation workloads, with initial prompt length T and number of generated tokens K . The throughput is computed as number of generated tokens over latency. For each setting (and additional benchmarks), we provide details in Appendix D.4.

Peak throughput Distilled models do not need kv -caches. This reduces memory requirement during generation, enabling higher peak throughput in large-batch workloads. We achieve 10 × higher throughput than Transformers at size 1 . 3 billion parameters (Figure 1.1). Throughput is higher than Transformers even at fixed batch sizes, indicating lower latency.

SSMstate dimension and throughput For typical distillation orders ( < 100 ), peak throughput is not greatly affected. We measure a 2% reduction in throughput from 32 to 64 .

Prompt length The throughput of LaughingHyena -distilled models is 4 × larger than Transformers at fixed batch size 64 and prompt length 1536 (Figure 5.3). As prompt length increases, the runtime gap between pre-filling via convolutions in LCSMs and pre-filling in Transformers widens (e.g., ˜ O ( T ) as detailed in Section 3.4, compared to O ( T 2 ) ).

Memory footprint Recurrent models do not require kv -caches and use constant memory for generation of an arbitrary number of tokens (Figure 5.4).

6 Conclusion

We study the efficiency and quality of state-of-the-art long convolutional sequence models. First, we introduce LaughingHyena , a novel distillation method inspired by rational function approximation and model-order reduction techniques. LaughingHyena can be applied after training to extract compact state-space models from each convolutional filter, without loss of quality. Distilled models achieve higher throughput than equivalently-sized Transformers, and can perform auto-regressive generation in constant memory by sidestepping the need to cache previous outputs. We theoretically and empirically investigate the trade-offs of different strategies for fast inference of recurrent models, and introduce architectural improvements to Hyena that improve pretraining quality.

512

1

,

024

1

,

536

1

,

000

2

,

000

3

,

000

Prompt Length

Throughput

[

tok

/s

]

Laughing Hyena 1

.

3B

Hyena 1

.

3B

Transformer 1

.

3B

H3 1

.

3B

64

256

512

10

20

30

Generation Length

Memory

[

GBs

]

Laughing Hyena 1

.

3B

Hyena 1

.

3B

Transformer 1

.

3B

H3 1

.

3B

Acknowledgments

We would like to thank Together Computer for providing the compute used to train models in this paper. We gratefully acknowledge the support of NIH under No. U54EB020405 (Mobilize), NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); US DEVCOM ARL under No. W911NF-21-2-0251 (Interactive Human-AI Teaming); ONR under No. N000141712266 (Unifying Weak Supervision); ONR N00014-20-1-2480: Understanding and Applying Non-Euclidean Geometry in Machine Learning; N000142012275 (NEPTUNE); NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASF, Accenture, Ericsson, Qualcomm, Analog Devices, Google Cloud, Salesforce, Total, the HAI-GCP Cloud Credits for Research program, the Stanford Data Science Initiative (SDSI), Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Program, and members of the Stanford DAWN project: Facebook, Google, and VMWare. This work is supported by NSF (1651565), AFOSR (FA95501910024), ARO (W911NF-21-1-0125), ONR, DOE (DE-SC0022222), CZ Biohub, and Sloan Fellowship. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of NIH, ONR, or the U.S. Government. AR's work is supported by NSF grant# CCF-2247014.

Broader Impact

In this work, we focus on advances related to efficient models for long sequences.

Efficiency Our distillation methods for constant-memory, high throughput inference in long convolution sequence models ( LCSMs ) can lead to energy savings during model deployement, enabling processing of longer-form content at a fraction of the cost and reducing environmental impact. Improved efficiency may also affect other aspects of AI safety, as it may make it easier produce malicious or harmful content.

Accessibility By improving the efficiency of training and generation, LCSMs and LaughingHyena may contribute to increased accessibility of large language models, lowering the hardware barrier to entry for individuals and organizations with limited resources.

Steerability New method based on LCSMs enable sequence models to process long-form prompts previously inaccessible by Transformers, which may lead to increased control over models via e.g., conditioning on additional instructions [45].

References

Daniel Y Fu et al. 'Hungry Hungry Hippos: Towards Language Modeling with State Space Models'. In: 2023 (cit. on pp. 1-3, 7, 23).

Michael Poli et al. 'Hyena Hierarchy: Towards Larger Convolutional Language Models'. In: (2023). arXiv: 2302.10866 (cit. on pp. 1-3, 7, 8, 31, 35, 43).

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. 'Neural machine translation by jointly learning to align and translate'. In: (2014). arXiv: 1409.0473 (cit. on p. 1).

Ashish Vaswani et al. 'Attention is all you need'. In: Advances in neural information processing systems 30 (2017) (cit. on p. 1).

Chi-Tsong Chen. Linear system theory and design . Saunders college publishing, 1984 (cit. on pp. 1, 6, 19, 21).

Albert Gu, Karan Goel, and Christopher Ré. 'Efficiently modeling long sequences with structured state spaces'. In: (2021). arXiv: 2111.00396 (cit. on pp. 1, 2, 4, 6, 20, 24).

Vadim Movsesovich Adamyan, Damir Zyamovich Arov, and Mark Grigor'evich Krein. 'Analytic properties of Schmidt pairs for a Hankel operator and the generalized Schur-Takagi problem'. In: Matematicheskii Sbornik 128.1 (1971), pp. 34-75 (cit. on pp. 2, 7, 29).

Yuji Nakatsukasa, Olivier Sète, and Lloyd N Trefethen. 'The AAA algorithm for rational approximation'. In: SIAM Journal on Scientific Computing 40.3 (2018), A1494-A1522 (cit. on pp. 2, 6).

GRB Prony. 'Essai experimental et analytique sur les lois de la dilatalrlite de fluids elastiques et sur cells de la vapeur de l'alcool, à différents tempoeratures'. In: Journal de l'Ecole Polytechnique (París) 1 (1795), pp. 24-76 (cit. on pp. 2, 6).

Alec Radford et al. 'Language models are unsupervised multitask learners'. In: OpenAI blog 1.8 (2019), p. 9 (cit. on p. 2).

Leo Gao et al. 'The pile: An 800gb dataset of diverse text for language modeling'. In: (2020). arXiv: 2101.00027 (cit. on pp. 2, 8, 35).

Albert Gu et al. 'Hippo: Recurrent memory with optimal polynomial projections'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 1474-1487 (cit. on pp. 2, 24).

David W Romero et al. 'Ckconv: Continuous kernel convolution for sequential data'. In: (2021). arXiv: 2102.02611 (cit. on p. 2).

David W Romero et al. 'Flexconv: Continuous kernel convolutions with differentiable kernel sizes'. In: (2021). arXiv: 2110.08059 (cit. on p. 2).

Vincent Sitzmann et al. 'Implicit neural representations with periodic activation functions'. In: Advances in neural information processing systems 33 (2020), pp. 7462-7473 (cit. on p. 3).

Rizal Fathony et al. 'Multiplicative filter networks'. In: International Conference on Learning Representations . 2020 (cit. on p. 3).

Daniel Y. Fu et al. 'Simple Hardware-Efficient Long Convolutions for Sequence Modeling'. In: International Conference on Machine Learning (2023) (cit. on pp. 3, 24).

Kemin Zhou and John Comstock Doyle. Essentials of robust control . Vol. 104. Prentice hall Upper Saddle River, NJ, 1998 (cit. on p. 4).

Athanasios C Antoulas. Approximation of large-scale dynamical systems . SIAM, 2005 (cit. on pp. 4, 7, 29).

Wilhelmus HA Schilders, Henk A Van der Vorst, and Joost Rommes. Model order reduction: theory, research aspects and applications . Vol. 13. Springer, 2008 (cit. on p. 4).

Sun-Yuan Kung. 'A new identification and model reduction algorithm via singular value decomposition'. In: Proc. 12th Asilomar Conf. on Circuits, Systems and Computer . 1978, pp. 705-714 (cit. on p. 4).

DFriedman. 'On approximating an FIR filter using discrete orthonormal exponentials'. In: IEEE Transactions on Acoustics, Speech, and Signal Processing 29.4 (1981), pp. 923-926 (cit. on p. 4).

J Bednar. 'On the approximation of FIR by IIR digital filters'. In: IEEE Transactions on Acoustics, Speech, and Signal Processing 31.1 (1983), pp. 28-34 (cit. on p. 4).

Bartlomiej Beliczynski, Izzet Kale, and Gerald D Cain. 'Approximation of FIR by IIR digital filters: An algorithm based on balanced model reduction'. In: IEEE Transactions on Signal Processing 40.3 (1992), pp. 532-542 (cit. on pp. 4, 44).

William Henry Young. 'On the multiplication of successions of Fourier constants'. In: Proceedings of the Royal Society of London. Series A, Containing Papers of a Mathematical and Physical Character 87.596 (1912), pp. 331-339 (cit. on p. 5).

William Beckner. 'Inequalities in Fourier analysis on Rn'. In: Proceedings of the National Academy of Sciences 72.2 (1975), pp. 638-641 (cit. on p. 5).

John Fournier. 'Sharpness in Young's inequality for convolution'. In: Pacific Journal of Mathematics 72.2 (1977), pp. 383-397 (cit. on p. 5).

Tong S Quek and Leonard YH Yap. 'Sharpness of Young's inequality for convolution'. In: Mathematica Scandinavica 53.2 (1983), pp. 221-237 (cit. on p. 5).

Jean-Paul Berrut and Lloyd N Trefethen. 'Barycentric lagrange interpolation'. In: SIAM review 46.3 (2004), pp. 501-517 (cit. on p. 6).

Marcos Vicente Moreira and João Carlos Basilio. 'Fair and Square Computation of Inverse ZTransforms of Rational Functions'. In: IEEE Transactions on Education 55.2 (2011), pp. 285290 (cit. on p. 6).

Henri Padé. 'Sur la représentation approchée d'une fonction par des fractions rationnelles'. In: Annales scientifiques de l'Ecole normale supérieure . Vol. 9. 1892, pp. 3-93 (cit. on pp. 6, 27).

Ankit Gupta, Albert Gu, and Jonathan Berant. 'Diagonal state spaces are as effective as structured state spaces'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 2298222994 (cit. on pp. 6, 24).

Antonio Orvieto et al. 'Resurrecting Recurrent Neural Networks for Long Sequences'. In: (2023). arXiv: 2303.06349 (cit. on pp. 6, 24, 26).

Jeffrey M Hokanson. 'A data-driven McMillan degree lower bound'. In: SIAM Journal on Scientific Computing 42.5 (2020), A3447-A3461 (cit. on p. 6).

L Ho and Rudolf E Kalman. 'Effective construction of linear state-variable models from input/output functions'. In: at-Automatisierungstechnik 14.1-12 (1966), pp. 545-548 (cit. on p. 6).

Carl Eckart and Gale Young. 'The approximation of one matrix by another of lower rank'. In: Psychometrika 1.3 (1936), pp. 211-218 (cit. on p. 7).

Guy E Blelloch. 'Prefix sums and their applications'. In: (1990) (cit. on p. 7).

Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. 'Simplified state space layers for sequence modeling'. In: (2022). arXiv: 2208.04933 (cit. on p. 7).

Noam Shazeer. 'Fast transformer decoding: One write-head is all you need'. In: (2019). arXiv: 1911.02150 (cit. on p. 8).

Joshua Ainslie et al. 'GQA: Training Generalized Multi-Query Transformer Models from MultiHead Checkpoints'. In: (2023). arXiv: 2305.13245 (cit. on p. 8).

Percy Liang et al. 'Holistic evaluation of language models'. In: (2022). arXiv: 2211.09110 (cit. on pp. 8, 9, 36).

Leo Gao et al. A framework for few-shot language model evaluation . Version v0.0.1. Sept. 2021. DOI: 10.5281/zenodo.5371628 . URL: https://doi.org/10.5281/zenodo.5371628 (cit. on pp. 8, 9, 36).

Ari Holtzman et al. 'The curious case of neural text degeneration'. In: (2019). arXiv: 1904. 09751 (cit. on p. 9).

Stella Biderman et al. 'Pythia: A suite for analyzing large language models across training and scaling'. In: (2023). arXiv: 2304.01373 (cit. on pp. 9, 36).

Yuntao Bai et al. 'Constitutional AI: Harmlessness from AI Feedback'. In: (2022). arXiv: 2212. 08073 (cit. on p. 11).

Alan V Oppenheim. Discrete-time signal processing . Pearson Education India, 1999 (cit. on p. 21).

Lennart Ljung. System identification . Springer, 1998 (cit. on p. 21).

RP Guidorzi. 'Certain models from uncertain data: the algebraic case'. In: Systems & control letters 17.6 (1991), pp. 415-424 (cit. on p. 21).

Irwin W Sandberg. 'On the theory of linear multi-loop feedback systems'. In: Bell System Technical Journal 42.2 (1963), pp. 355-382 (cit. on p. 22).

Michael Zhang et al. 'Effectively Modeling Time Series with Simple Discrete State Spaces'. In: (2023). arXiv: 2303.09489 (cit. on p. 24).

Luca Perotti and Michał Wojtylak. 'Matrix methods for Padé approximation: Numerical calculation of poles, zeros and residues'. In: Linear Algebra and its Applications 548 (2018), pp. 95-122 (cit. on p. 27).

Amer Abu-Omar and Fuad Kittaneh. 'Estimates for the numerical radius and the spectral radius of the Frobenius companion matrix and bounds for the zeros of polynomials'. In: Annals of Functional Analysis 5.1 (2014), pp. 56-62 (cit. on p. 27).

Julia Eaton et al. 'Polynomial root radius optimization with affine constraints'. In: Mathematical Programming 165 (2017), pp. 509-528 (cit. on p. 27).

Gerlind Plonka and Vlada Pototskaia. 'Application of the AAK theory for sparse approximation of exponential sums'. In: (2016). arXiv: 1609.09603 (cit. on p. 29).

Jimmy Ba et al. 'Using fast weights to attend to the recent past'. In: Advances in neural information processing systems 29 (2016) (cit. on p. 31).

William B Johnson. 'Extensions of Lipschitz mappings into a Hilbert space'. In: Contemp. Math. 26 (1984), pp. 189-206 (cit. on p. 31).

Devdatt P Dubhashi and Alessandro Panconesi. Concentration of measure for the analysis of randomized algorithms . Cambridge University Press, 2009 (cit. on p. 34).

Ilya Loshchilov and Frank Hutter. 'Decoupled weight decay regularization'. In: (2017). arXiv: 1711.05101 (cit. on p. 35).

Dale F. Enns. 'Model reduction with balanced realizations: An error bound and a frequency weighted generalization'. In: The 23rd IEEE Conference on Decision and Control (1984), pp. 127-132 (cit. on p. 44).

Contents

LAUGHING HYENA DISTILLERY

Supplementary Material

Stefano Massaroli, Michael Poli, and Dan Fu contributed equally to this work. Christopher Ré, Stefano Ermon, and Yoshua Bengio share equal senior authorship.

All authors read and approved the final manuscript.

A Linear Systems

A.1 Extended Notation and System Theory Preliminaries

We first introduce the notation and some mathematical concepts that will be used throughout the paper. By Z we denote the set of integers, by R the set of reals, and by C the set of complex numbers. The variable t stands for time . ℓ p ( Z ) denotes the Banach space of complex-valued sequences ( x t ) t ∈ Z with finite energy, i.e. ∥ x ∥ p := [ ∑ t ∈ Z | x t | p ] 1 /p < ∞ for some 1 ≤ p < ∞ . ℓ ∞ ( Z ) is instead is the space of sequences for which ∥ x ∥ ∞ := sup t ∈ Z | x t | < ∞ . With S denoting the unit circle in the complex plane, S := { z ∈ C : | z | = 1 } we define H p ( S ) as the space of functions X from C to itself such that ∥ X ∥ p := [(1 / 2 π ) ∫ π -π | X ( e iω ) | p d ω ] 1 /p < ∞ and H ∞ ( S ) the space for which ∥ X ∥ ∞ := sup z ∈ S | X ( z ) | < ∞ . Particularly, K 2 ( S ) is a Hilbert space with inner product ⟨ X,Y ⟩ := (1 / 2 π ) ∫ π -π X ( e iω ) Y ∗ ( e iω )d ω where ' ∗ ' denotes complex conjugation. Although we acknowledge we are using the same notation for norms in both ℓ p ( Z ) and H p ( S ) , the correct meaning will always be made clear by the context. The Z -transform of a sequence x = ( x t ) t ∈ Z is X ( z ) = Z [ x ]( z ) := ∑ t ∈ Z x t z -t . We embrace the system theory convention of using capital letters to identify transformed sequences. The Z -transform is a projection of the sequence onto a basis of powers e t = r -t e iωt . This basis is not orthogonal unless r = 1 . That is the basis of the discrete-time Fourier transform F . Hence, F is defined as F [ x ]( e iω ) = X ( e iω ) := ∑ t ∈ Z x t e -iωt . The discrete-time Fourier transform is an isometric isomorphism between ℓ 2 ( Z ) and L 2 ( S ) . We say that sequences live in the time domain and their Z (or F ) transforms in the frequency domain .

A linear system is a linear operator transforming an input sequence u to an output sequence y . If the sequences have continuous support, i.e. t ranges over a continuous set (e.g. R ), we have a continuoustime system. Conversely, if the sequences have discrete support, i.e. t ranges over a discrete set (e.g. Z ), we have a discrete-time or digital system. In this manuscript we restrict ourselves to discrete-time systems . Systems can be single-input single-output (SISO) if u and y are scalar functions or multiinput multi-output if either u or y are vector-valued. We limit our discussion to SISO systems . The impulse response of a system is the output sequence y when the input sequence u is the Kronecker delta function δ t and is usually denoted by the letter h . The values h t of the impulse response sequence are also known as the Markov parameters of the system. The most common mathematical representation of a linear system is its convolution form: y = h ∗ u , i.e. y t = ∑ j ∈ Z h t -j u j = ∑ j ∈ Z h j u t -j , t ∈ Z . In matrix form the input-output relation is given by the Toeplitz operator T h corresponding to the (possibly infinitely long) sequence h , i.e. y = T h u . Taking the Z -transforms, we can write the input-output relation as Y ( z ) = H ( z ) U ( z ) (this is just the Fourier convolution theorem extended outside the unit circle). H ( z ) is called the transfer function of the system. When z = e iω , H ( e iω ) is just the discretetime Fourier transform of h which is called the frequency response of the system. A linear system is causal if h t = 0 for t < 0 . A system is called stable if the T h is a bounded operator. If u, y ∈ ℓ 2 ( Z ) , then stability implies h ∈ ℓ ∞ . In the following, we mainly focus on causal stable systems .

A.2 Systems Norms

When quantitatively characterizing linear systems, several norms play a crucial role. These norms provide measures of various characteristics of the systems, which are essential in both analysis and filter design.

The ℓ 2 and H 2 norms As defined above, the ℓ 2 norm represents the energy of a signal h , while H 2 is the energy of the (continuous) spectrum of h ,

By Parseval's theorem, the ℓ 2 and H 2 norms are equal, ∥ h ∥ 2 = ∥ H ∥ 2 . Further these norms are useful to study the approximation of convolutional filter. The following holds:

Lemma A.1 ( ℓ ∞ output error) . Consider the class of ℓ 2 measurable inputs such that ∥ u ∥ 2 ≤ ζ , then for all H, ˆ H ∈ H 2 ,

Proof.

If u is the unit impulse function u t = δ t then ζ = 1 . The results also holds for finite sequences of length L using the discrete Fourier transform.

Lemma A.2 (Impulse response error on finite sequences) . Consider filters h, ˆ h with finite length L . Then, the following holds.

where H and ˆ H denote the discrete Fourier transforms of h and ˆ h , respectively.

Proof.

A.3 Transfer Function of State-Space Models

The transfer function (3.1) is derived by taking the z -transform of input and state, U ( z ) = Z [ u ]( z ) , X ( z ) = Z [ x ]( z ) . Plugging U ( z ) , X ( z ) in the state equation (2.2), it holds

Substituting in the output equation yields

The transfer function is then defined as

Alternative derivation The transfer function can also be derived by direct z -transform of the impulse response h t of the system. This derivation is useful to highlight the region of convergence of the transfer

function.

We look at the convergence of the series ∑ ∞ t =0 ∥ z -1 A ∥ t 2 . We have

The series converges to 1 / (1 -r -1 ρ ( A )) if and only if r -1 ρ ( A ) < 1 i.e. for r > ρ ( A ) . Thus, in the exterior of the disk with radius ρ ( A ) , D ρ ( A ) := { z ∈ C : | z | > ρ ( A ) } , ∑ ∞ t =0 ( z -1 A ) t converges to ( I -z -1 A ) -1 and

The transfer function H ( z ) = h 0 + C ( z I -A ) -1 B of a stable lumped discrete-time system is defined outside the disc in the complex plane that encloses all the eigenvalues of A .

Invariance of the transfer function H ( z ) as defined in (A.1) is a proper 13 rational function of z . In case h 0 = 0 , H ( z ) is strictly proper and the denominator is monic:

Specifically, the denominator could be derived from A with det ( z I -A ) , and the numerator is det ( z I -A + BC ) + det ( z I -A ) . We provide a detailed derivation below in Section A.6. While state-space representation involves the analysis and synthesis of model matrices A , B , C , the transfer function is entirely characterized by the coefficients a = ( a n ) d n =1 , b = ( b n ) d n =1 of numerator and denominator polynomials. Notably, the transfer function is an invariant of the system: if we apply a change of variables to the state, the transfer function remains unchanged.

Lemma A.3. Coefficients a, b are invariant under any invertible change of variables.

Proof. The proof can be found in [5, pp.95] and follows from the definition of equivalence transformation . Consider the state-space matrices of under change of variables ˆ x = K x ,

The resulting transfer function H ( z ) can then be computed as

A.4 Truncated Transfer Functions

In the case of generic truncated (finite) impulse response filters, such that h t = 0 for all t greater than a certain value L (which we refer to as the length of the filter), the transfer function is simply a polynomial in the complex variable z of order L , i.e.

In case the filter is generated by a finite dimensional (lumped parameters) system, i.e. h t = CA t -1 B t = 1 , . . . , L , then (A.4) can still be represented exactly by a rational function of order d .

Lemma A.4 (Truncated rational transfer functions) . Consider the L -truncated impulse response h t ∈ ℓ 2 ( N ) of a lumped-parameter filter ( A , B , C , h 0 ) ,

Then its truncated transfer function is

Proof. By definition of z -transform we have

The sum ∑ L -1 t =0 ( z -1 A ) t is a partial Neumann series and can be manipulated as follows.

Thus,

which plugged in (A.5) gives H L ( z ) = h 0 + C ( I -z -L A L )( z I -A ) -1 B , proving the result.

Because of truncation, evaluating the transfer function H L ( z ) on the L roots of unity z = e iω k , w k = 2 πk/T for k = 0 , . . . L gives the lengthL discrete Fourier transform (DFT) of the filter:

In practice, this means that ¯ H ∈ C L is the FFT of h , ¯ H = FFT L [ h ] . If we can find an efficient and stable algorithm to evaluate ¯ H from the system matrices ( A , B , C , h 0 ) , then the FFT -based convolution of truncated filter with an input sequence u ∈ R L can be evaluated in ˜ O ( L ) time.

Reparametrization Assume training a LCSM equipped with SSM filters with input/target sequences to be all of length L (smaller sequences can be padded with zeros to the maximum length). Thus, for training purposes, we are only interested in evaluating ¯ H for the FFT -based convolution.

The truncated transfer function H L is equal to the original one with a correction term I -z -L A L on the numerator polynomial. As already noted in S4 [6], z -L is conveniently equal to one on the roots of unity, z iω k L = e -i 2 πk = 1 for all k = 0 , . . . , L -1 . Hence, the correction term due to truncation becomes constant: H k = C ( I -A L )(exp( -i 2 πk/L ) I -A ) -1 B ; in DFT domain the truncated filter behaves as the infinitely long one with a perturbed C matrix

If -as assumed- the SSM is stable ρ ( A ) < 1 , ( i ) the transfer function is defined on the unit circle, term CA L will go to zero exponentially fast as L → ∞ and ¯ C = C (as expected). As advised in [6], it is desirable to parametrize directly ¯ C ; the expensive computation of the correction term C ( I -A L ) is never carried out during training. Instead, the real C matrix can be retrieved for recurrent inference by inverting the correction term C = ¯ C ( I -A L ) -1 , always invertible for stable systems although possibly ill conditioned by eigenvalues too close to the stability margin (the unit circle).

A.5 From Transfer Function to State-Space

Suppose the coefficients of the numerator and denominator polynomials of a proper transfer function H is given:

A state-space representation of the form (2.2) can be rapidly realized in two steps:

and

We then use the long division method to compute b 0 /p ( z ) :

to finally get

Get delay-free path From (A.6) we first notice that the bias term h 0 is h 0 = b 0 . We thus want to isolate b 0 from the rest of the numerator. This can be obtained via long division (see §A.5.1) and results in

Get state-space matrices Given the transfer function H ( z ) with the isolated pass-through coeffient b 0 as in (A.7), we can construct the state-space matrices by companion canonical realization:

Details on the complete a la [5] derivation can be found in §A.5.2. A linear system with finitedimensional state can be equivalently characterized: by its state-space matrices ( A , B , C , h 0 ) , by its impulse response function h , or by the coefficients a, b (or β ) of the transfer function. A fourth representation is its linear-constant-coefficients difference equation form

typically used in signal processing literature in the theory of infinite impulse response filters (see [46]) and known, in the context of system identification of error-in-variables models, as auto-regressive moving-average filters [47, 48].

A.5.1 Isolating the h 0 -term from Transfer Function by Long division

If the rational transfer function H ( z ) accounts for the h 0 term, then it is simply proper (order of numerator equals the order of denominator), h 0 is necessarily h 0 = b 0 (the delay-free path). Given the transfer function in this form, we can isolate the b 0 term and the strictly rational term of (A.3) by long division. We start by expanding the fraction as

· · · Note that the coefficients b n in (A.3) correspond to b n -b 0 a n in (A.6), b n ← b n -b 0 a n . It is indifferent to parameterize the coefficients of the transfer function in either forms. However, if we choose the simply proper representation (A.6), we need to apply the derived correction factor to the numerator coefficients when we separate the h 0 term and strictly proper part of H ( z ) .

A.5.2 Construction of the State-Space from the Transfer Function

Chen's derivation The derivation is based on the steps reported for the continuous-time multi-input multi-output case in [5]. First, we define a pseudo-state v such that

Then, we define the state x t := ( x 1 t , . . . , x d t ) ∈ R d as

From (A.9) we have

Thus, we have the overall recurrence

which can be written in matrix form as

The output spectrum is then given by

Therefore, and the output equation in time-domain is given by

yielding state-space matrices (A.8).

A.6 From State-Space to Transfer Function

We detail an implementation oriented method to compute the coefficients ( a n ) d n =1 , ( b n ) d n =0 of a SSM's transfer function. Recall that

Hence, the denominator coefficients ( a n ) d n =1 are simply the coefficients of the characteristic polynomial of matrix A . They can be easily obtained by 1. computing the eigenvalues of A and 2. calculating the coefficients of the polynomial whose roots are such eigenvalues. On the other hand, the numerator apparently involves more complex symbolic manipulation. This can be simplified recalling a classic matrix-determinant identity:

Lemma A.5 ([49]) . Let M , B , and C respectively denote matrices of orders d × d , d × 1 , and 1 × d . Then,

Applying Lemma A.5 to (A.11) we obtain

Let poly ( r ) denote the coefficients of the polynomials with roots r = ( r 1 , . . . , r d ) . Then a = poly ( eig ( A )) . Since A and A -BC are of equal dimension, their characteristic polynomials have equal order and therefore

A.7 State-Space Representation of Truncated Filters.

A truncated filter h 0 , . . . , h L - as the ones found in any standard convolutional neural network - can be represented by a L -dimensional companion canonical SSM. The filter's transfer function H ( z ) = h 0 + h 1 z -1 + · · · + h L z -L is polynomial, i.e. a rational function with the denominator's coefficients set to zero. Following the canonical realization process detailed in Section A.5, the truncated filter has state-space form:

t

y

= [

h

1

h

2

h

L

]

x

t

+

h

0

u

t

.

· · · If x 0 = 0 L and u t = 0 for negative t , then at each t > 0 the state is a shifted copy of the input sequence x t = ( u t -1 , . . . , u t -L ) ∈ R L . Nonetheless, the asymptotic complexity of computing one recurrent step is O ( L ) as it requires only a shift operation and a lengthL dot product

The memory footprint is also O ( L ) . In [1] it is proposed the use of shift-type SSMs to parametrize one of the filters of the H3 block.

A.8 Efficient Computation of State-Space Models

A.8.1 Fast Evaluation of the Transfer Function

Computing H ( z ) at any point z ∈ C concerns the evaluation of the d -order polynomial of numerator and denominator,

In practice, we are mainly interested in a fast algorithm that allows computing H on the L roots of unity to obtain the DFT of the filter. The DFT of the filter can be then readily used to perform a FFT -based convolution with a lengthL input sequence u or to recover the impulse response function via inverse DFT. We prove the following:

Lemma A.6. Given the coefficients a , b of the transfer function, the frequency and impulse response of the filter can be evaluated in ˜ O ( L ) time.

Proof. The result is proven showing that the transfer function can be evaluated in ˜ O ( L ) time on the L roots of unity. The fastest method to evaluate polynomials on L arbitrary points z of the complex plane is generally the Horner's scheme. This method is based on a sequence of nested multiplications and computes the polynomial from its vector of coefficients, delivering a time complexity of O ( dL ) . More explicitly, Horner's scheme determines p ( z ) as p ( z ) = (( · · · (( a d z -1 + a d -1 ) z -1 + a d -2 ) · · · ) z -1 + a 2 ) z -1 + a 1 ) z -1 +1 . Each step involves a multiplication and an addition, making a total of 2 d operations per evaluation point. Thus, for L points, the total number of operations amounts to O ( dL ) .

Effectively, Horner's approach implements the matrix-vector product of an L -by-( d +1) Vandermonde matrix V ∈ C L × ( d +1) constructed by L evaluation points ( z 0 , . . . , z L -1 ) with the vector of coefficients

Significantly, if the polynomial is required to be evaluated at the roots of unity, the Vandermonde matrix simplifies corresponds to the L × ( d +1) DFT matrix. Further, zero-padding the coefficient vector to length L , enables the use a single lengthL FFT to compute the matrix-vector product in ˜ O ( L ) time. Thus, the numerator and denominator polynomials of the transfer function can be evaluated, on the roots of unity, in ˜ O ( L ) time by taking the FFT of the padded numerator / denominator coefficients a, b and subsequently dividing element-wise the two sequences as FFT L [ b ] / FFT L [ a ] . The overall time complexity to obtain the impulse response is also ˜ O ( L ) since h can be recovered taking an inverse FFT of the frequency response.

A.8.2 Fast Companion Recurrence

The recurrent step of a generic SSM (2.2) with dense system matrices usually requires O ( d 2 ) operations due to the matrix-vector product A x t . We show how the recurrence of SSMs in companion canonical form , i.e. with system's matrices (A.8), requires only O ( d ) operations.

Lemma A.7. The recurrent step of a state-space model in companion canonical form (A.8) can be evaluated in O ( d ) time and memory.

Proof. The companion state matrix A can be broken down into a lower shift matrix L N and a low-rank term. Particularly, with e 1 the first element of the canonical basis of R N and α = ( a 1 , . . . , a N ) , we have

It follows that the recurrent update can be simplified to

The peculiarity of this formulation is that we never need to construct the matrices to perform the recurrence. In particular we have:

Thus, each step only requires two inner products ( d multiplications and d sums each) and one shift operation, totaling O ( d ) operations.

The proof of Lemma A.7 yields the practical implementation of the recurrence:

A.8.3 Canonization of State-Space Models

The companion canonical form discussed in Section A.5 is the ideal representation to deploy SSM-based convolutional layers: i ) it comes with a O ( d ) fast recurrence and ii ) allows to swiftly switch between time and frequency domains with a direct mapping between state-space matrices and coefficient of the transfer function (which in turn allow ˜ O ( L ) fast convolutions).

Aside from [50], which directly parametrizes S4 layers in companion canonical form, all the other parameterizations [12, 6, 32, 17, 33] can be converted ( canonized ), under mild assumptions.

Lemma A.8 (Canonization of SSMs) . Any state-space model (2.2) with proper transfer function can be converted in companion canonical form.

Proof. The result can be proved following the two-step conversion process.

Get the coefficients of the transfer function : Given the original state-space matrices ( A , B , C , h 0 ) , the transfer function is given by H ( z ) = C ( z I -A ) -1 B + h 0 . A proper rational function has the form H ( z ) = q ( z ) /p ( z ) where the numerator q ( z ) has coefficients b = ( b n ) d n =0 and the denominator has coefficients a = ( a n ) d n =0 ( a 0 = 1 since p is monic). As shown in Section A.6, the coefficients of the transfer function can be extracted in closed-form as b = poly ( eig ( A -BC )) + poly ( eig ( A ))(1 -h 0 ) and a = poly ( eig ( A )) 14 ;

Construct companion matrices Given the coefficients a and b a new set of canonical statespace matrices which realize the transfer function can be obtained following the recipe of Section A.5.

The resulting companion SSM is equivalent the the original one since they share the same transfer function.

B LaughingHyena : Further Details

B.1 Parametrization of Modal Interpolators

Complex-conjugate states Assuming even distilling dimension d , we pick poles λ n and residues R n in complex-conjugate pairs:

which allow partitioning the state-space matrices as

where

If we also partition the state as x = (¯ x, ˜ x ) , ¯ x, ˜ x ∈ C d/ 2 , the resulting recurrence has the form

We have

Thus, if ˜ x 0 = ¯ x ∗ 0 , then ˜ x t = ¯ x ∗ t for all t > 0 . Hence, at inference time we only need to propagate forward half of the state - say ¯ x - and then compute the output as

This parametrization allows to update only half of the state, reducing the time and memory cost compared to an unconstrained linear system with complex coefficients. However, the implicitly achieved realness of the output (assuming D = h 0 ∈ R and u t ∈ R ) comes at a cost of expressivity: such a system is equivalent to an unconstrained complex linear system of dimension d/ 2 of which we only keep the real part of the output.

Poles and residues For the modal interpolation, the parametrization is analogous to the one of a diagonal state space model [33]. Poles λ n and residues R n need both to be complex numbers. In [33] the authors suggest parametrizing real and imaginary components of B and C matrices while representing the eigenvalues λ n in polar form, λ n = r n e iα n with r n and α n being themselves exponential functions of the actual trainable parameters, r n = e -e νn , α n = e ζ n leading to

This ensures stability of the poles | λ n | < 1 and positive-only phases α n . For the purpose of distillation we propose a simplified parametrization as follows:

We only parametrize the C vector. Parametrizing both B and C is redundant and increases the computational cost of performing each step of the recurrence. The residues R n correspond in fact to R n = C n B n of a diagonal state space model. Setting B = 1 d saves parameters without harming expressivity. Further if B is different from 1 k it needs to be multiplied to u t at each recurrence step. C n = R n = R [ R n ] + i I [ R n ] and R [ R n ] , I [ R n ] are the trainable parameters of the residue.

For the purpose of distillation we have no benefit in forcing the eigenvalues of the model to be stable, i.e. constrained to lie strictly inside the unit circle. Instead, such constrain may actually harm the expressivity of the approximant. We choose the the simpler parametrization λ n = r n e iα n , r n , α n ∈ R .

B.2 Distillation as Rational Interpolation

Distillation as rational interpolation Approximating a filter with an SSM can be thus achieved by fitting a proper rational function to the (truncated) transfer function of the original filter H L ( z ):= ∑ L t =0 h t z -t . That is,

Amodern 15 way to solve this problem by H 2 error minimization via gradient descent 16 . We can use the Fast Fourier Transform ( FFT ) to evaluate both the target and distilled transfer functions and solve:

To ensure stability of the distilled filters and well-conditioned gradient descent dynamics, the roots of the denominator polynomial must strictly lie inside the unit circle ( ρ ( A ) < 1 ). This, in turn, requires constraining the coefficients a into the region { a : poly ( a ) is stable } which is by itself an open research problem [52, 53]. Experimentally, we observe that standard coefficient normalization techniques overly restrict the parameters space and lead to poor distillation performances at reasonable order.

C Proofs

C.1 Proof of Lemma 2.1

Generating K tokens with a long convolution layer (2.1) from a lengthT prompt has time complexity O ( T log 2 T + TK + K 2 ) and requires O ( L ) memory.

Proof. We compute the time complexity memory of a lengthT prompt processing ( pre-filling ) and subsequent auto-regressive decoding of K tokens. The auto-regressive generation of long convolution computes the next token as by

The pre-filling step is needed to prime this recurrence by computing the first T outputs till y T -1 from the lengthT prompt u . This is just a convolution between two lengthT signal and requires O ( T log 2 T ) time and linear memory. The auto-regressive decoding of K tokens requires K steps (C.1) with the length of the sequences increasing by 1 at each step. Thus we have a total asymptotic complexity of

and requires at worst ( k = K -1 ) to store the length T + K = L generated output sequence, i.e. O ( L ) memory. In the limit we thus have a total time complexity of O ( T log 2 T + TK + K 2 ) and O ( L ) memory.

C.2 Proof of Lemma 2.2

Generating K tokens with a SSM (2.2) from a lengthT prompt has time complexity O ( T log 2 T + dK ) and requires O ( d ) memory.

Proof. In autoregressive mode, the cost of generating one token is the cost of evaluating the state recurrence (2.2). Each step then requires O ( d ) time and memory for the class of SSMs considered in this work (see Lemma A.8). Hence, generating K tokens costs O ( dK ) time and constant O ( d ) memory (we only need to store the current state).

The recurrence is initialized for autoregressive generation with the post-prompt state x T -1 and output y T -1 . The latter can be recovered in linear time and memory O ( T ) by definition y T -1 = ∑ T -1 j =0 h t -j u j (assuming to have the impulse response h available) and state x T -1 in O ( dT ) time and d memory through the recurrence. The overall asymptotic cost is therefore O ( dL ) time and O ( d ) memory.

Note that, for prompts and SSMs of practical sizes we usually have d > log 2 T . In such a case the state x T -1 can be computed in T log 2 T time rather than dT by Proposition 3.2.

C.3 Proof of Lemma 2.3

Generating K tokens with self-attention from a lengthT prompt has time complexity O ( T 2 + TK + K 2 ) and requires O ( L ) memory.

Proof. The proof is identical to the one of Lemma 2.1, with the only difference of a quadratic asymptotic cost O ( T 2 ) to process the prompt obtain the kv cache.

Self-attention suffers with long contexts: it is significantly more expensive in prefilling than long convolutions and SSMs due to its quadratic cost. Nonetheless, in autoregressive mode, self-attention reaches the same overall asymptotic complexity O ( TK + K 2 ) as long convolutions (with the memory overhead of having to cache k and v ).

C.4 Proof of Proposition 3.1

If A has semi-simple eigenvalues λ n ∈ C , then the transfer function of the system can be decomposed as ˆ H ( z )= ∑ d n =1 R n / ( z -λ n ) where R n ∈ C is the residue associated with the pole λ n .

Proof. If A is semi-simple, then it is diagonalized by a basis V of eigenvectors; it admits an eigenvalue decomposition diag ( λ ) = VAV -1 where λ = ( λ 1 , . . . , λ d ) ∈ C d contains the eigenvalues of A . Projecting the state onto the basis of eigenvectors, s := V x , the state space model is is transformed into modal form:

where ˜ B := VB = ( ˜ b n ) d n =1 and ˜ C := CV -1 = (˜ c n ) d n =1 . In modal form, the state equations are decoupled, i.e.

Letting R n = ˜ b n ˜ c n , proves the result.

C.5 Proof of Lemma 3.1

The distilled filter ˆ h in modal form (3.2) can be computed in O ( dL ) time from its modal form and in ˜ O ( L ) from its rational form.

Recalling (3.2), ˆ h t = ∑ d n =1 R n λ t -1 n , R n , λ n ∈ C , t > 0 , the O ( dL ) complexity of the impulse response is apparent: for each of the t = 1 , . . . , L , ˆ h t can be computed in O ( d ) time. The ˜ O ( L ) cost from the rational form follows by Lemma A.6.

C.6 Proof of Theorem 3.2

Let h be a lengthL filter, ˆ h a distilled filter of order d < L and let S L , ˆ S L be the respective Hankel matrices. Then inf ˆ S L ∥ S L -ˆ S L ∥ 2 = σ d .

Proof. The theorem characterizes the best-case scenario in terms of approximation error of the distilled SSMs or a certain order d where it is clear that rank ˆ S ≤ d . This theorem is a direct application of the Adamyan-Arov-Krein (AAK) theory of infinite Hankel operators [7]. Let ˆ S ∗ L = arg inf ˆ S L ∥ S L -ˆ S L ∥ 2 ; the AAK theorem says that every causal system can be optimally approximated by another causal system of lower dimension. Optimal here means

where the first infimum is taken over all Hankel matrices S ∗ L and the second over all arbitrary matrices K (see [19, Chapter 8] and [54] for further details and references).

C.7 Proof of Proposition 3.3

The filter (3.2) has a state space matrices A = diag ( λ 1 , . . . , λ d ) ∈ C d × d , B = (1 , . . . , 1) ⊤ ∈ C d × 1 , C = ( R 1 , . . . , R d ) ∈ C 1 × d , D = h 0 whose step can be evaluated in O ( d ) time and memory.

Taking the z -transform of the output equation and each state equation yields

Thus, the overall transfer function is

D is set to h 0 by default. The result is proven showing that (3.2) can be written in the form ˆ h t = CA t -1 B for t > 0 . If we choose A = diag ( λ ) then the impulse response becomes

The choice B n = 1 for all n = 1 , . . . , d , B = 1 d and C n = R n finalizes a modal canonical statespace realization of the distilled filter. The O ( d ) time complexity of the corresponding recurrent step is guaranteed by the decoupling of each state equation from another,

Each of the d state equations can be computed (in parallel) in O (1) time. The output equation is a dot product requiring d multiplications and d additions, hence the O ( d ) time compexity of the recurrence.

C.8 Proof of Proposition 3.2

x T = ( v T , . . . , v T -d ) where v = g ∗ u and g is the filter whose transfer function is 1 / den ( ˆ H )( z ) and can be evaluated in ˜ O ( T ) .

Without loss of generality, let us assume to have converted the distilled filter in canonical form (i.e. we have unrestricted access to the coefficients of the rational transfer function) and let D = 0 . We use the notation of Section A.5. In z -domain, the state-to-input relation is given by

On the other hand Y ( z ) = ˆ H ( z ) U ( z ) = q ( z ) /p ( z ) U ( z ) . Therefore,

Let V ( z ) = U ( z ) /p ( z ) . From the shift property of the z -transform it holds,

v can be obtained in ˜ O ( L ) time via an FFT -convolution of the input u and g , the filter resulting from inverse transforming 1 /p ( z ) . The proof is convoluded setting t = L

C.9 Proof of Theorem 4.1

Notation. Wewill be denoting the all 1 row vector of size k , given by [1 1 . . . 1 1] , and the all 0 row vector of size k , given by [0 0 . . . 0 0] , as 1 k and 0 k , respectively. We will also construe the standard basis vector e i as a column vector in these notes. Next, we will adhere to the following matrix indexing convention: A ij is the entry in the i th row and the j th column, A [ i, :] ∈ F 1 × n denotes the i th row, and A [: , j ] ∈ F m × 1 denotes the j th column of A ∈ F m × n . Here, we also use 0 m × n ∈ R m × n and I n to denote the matrix of all zeros and the identity matrix of dimension n , respectively. Moreover, we extend the outer product between two vectors to a tensor product using the symbol ⊗ , the computation of which is carried out batch-wise with some dimension of one or both of the input tensors. Finally, we express the binary encoding of i ∈ [ n ] in a row vector form, given by B i ∈ Z P n 2 , where P n is the closest power of 2 to n.

Language and Model Description. The language Λ has s keys and s values: L K := { k 1 , . . . , k s } , L V := { v 1 , . . . , v s } . Formally, the language Λ consists of sequences x ∈ ( L K × L V ) s × L K , where there is an associated mapping f x : L K → L V . For each sequence, the odd indices in [ L ] belong to L K , for x 1 , x 3 , . . . , x L , and we define

The last item x L ∈ { x 1 , x 3 , . . . , x L -1 } , called the query , must be one of the keys that has appeared in x already. Our goal is to produce f x ( x L ) at the end of the sequence, which we refer as the associated value . This problem is termed as the associative recall problem [55].

We will now outline the Hyena layer [2] with multiple heads as follows.

Algorithm 1 Hyena

Require: Input sequence u ∈ R L × D from the previous layer, long convolution filter T h , number of heads M .

2: for m = 1 , . . . , M do

1: q m , k m , v m ← Projection ( u ) for m ∈ [ M ] .

3: Perform the outer product z m ← k m ⊗ v m ∈ R L × N × N , where N := D/M .

5: Average the output y ← ( ∑ m ) y m /M

4: Apply the convolution independently and compute y t m ← T h ( z t m ) q t m ∈ R L × N

6: Retrieve the value f ( k L ) of the key k L from y [ L, :] .

In order to prove Theorem 4.1, we need the following technical statement concerning sparse recovery of a heavy-hitter.

Proposition C.1 (Heavy-Hitter Recovery) . Let x ∈ R s be a vector with one entry bounded by 1 ± 1 3 4 √ s -referred as the heavy-hitter -and the rest of the entries bounded by ± 1 3 4 √ s . Then, there exists a matrix S ( m ) ∈ R s × O ( √ s log s ) such that the position of the heavy-hitter in x can be inferred from the average of M measurements with S ( m ) given by (∑ m x S ( m ) ) /M with probability of error ≤ 1 s .

Before presenting the proof of Proposition C.1, we use it to prove Theorem 4.1 as follows.

Proof of Theorem 4.1. We take D = O ( √ s log 2 s ) and M = 243 · log s so that N = O ( √ s log s ) and use the same projections and filters for each head. We will start by describing the projections of the input. To this end, let E : [ L ] → 2 s define a map from the row indices of u to the keys k i and values f x ( k i ) given by

Here, we note that we also have

as the even indices are defined as x t = f x ( x t -1 ) for t even (C.3), whence x t -1 ∈ L K as t -1 is odd. Next, we can separate the keys q , queries q and values v from the input sequence u . For keys and queries, we will be using the Johnson-Lindenstrauss embedding [56]. We state its guarantee here.

For a set of points P ⊆ R s , let ϵ, δ > 0 with k ≥ 2 ln ( 2 s δ ) /ϵ 2 , and f : R s → R k be the randomly constructed linear map from [56], then with probability of error ≤ δ, we have

for all x, y ∈ P .

More precisely, we take R ∈ R O ( √ s log s ) × s to be the matrix representation of f with ϵ := 1 3 4 √ s and δ := 1 s c for some c > 1 so that R [: , i ] = f ( e i ) . Thus, we define

For values, we use the heavy-hitter recovery matrix as described in Proposition C.1 so that we have

Further, using 1DConv (equivalently, in terms of polynomials, h ( X ) := X ), we can shift the queries to get the projection for keys k so that we have k m [ t, :] = q m [ t -1 , :] The Hyena filters, along with the specific convolution being performed by T h , are specifically described in terms of polynomial multiplications, for all m ∈ 1 , . . . , M , as follows.

Here, we note that T h ( u ) takes the cumulative sum over the input. That is, for all i , we have

We will now compute z m as follows

Further, applying the convolution, we get

For inference, it suffices to show that the last row of the output y recovers the output with high probability. Indeed, let t ′ ∈ [ L ] denote the row index of the value associated to the query such that the corresponding key has the following relation

Finally, we multiply by the query q across L . Specifically, we now look at the computation of the L th row of y :

Here, we are using the fact that k m [ t ′ -1 , :] = q m [ L, :] due to (C.8) in (C.9). We then change the indexing from (C.9) and (C.11) by observing that all the odd entries corresponding to values are zeroed out in K (cf. C.6). Finally, we simply substitute (C.6) and (C.7) in (C.10) and (C.11), respectively. Next, we define x ∈ R s +1 with where j = t -1 with t ∈ [ L ] and t even. Here, x ∈ R n +1 is a vector of size s +1 as there are s +1 such even numbers in [ L ] . Note that x is the vector with a heavy-hitter from Proposition C.1. To see

̸

this, observe that we have ∣ ∣ x E ( t ′ -1) -1 ∣ ∣ ≤ 1 3 4 √ s and | x j | ≤ 1 3 4 √ s for all j = E ( t ′ -1) . Using (C.11), with probability ≥ 1 -1 s c , we then have

By Proposition C.1, we can then infer the position of the key at t ′ with probability of error 1 s . By the union bound, we can then retrieve the corresponding value with probability at least 1 -( 1 s + 1 s c ) .

We will now prove C.1 as follows.

Proof of C.1. We will assume that s is a power of 2 for the sake of simplicity. We first specify how we will construct such an S ( m ) ∈ R s × O ( √ s log s ) . Let h : [ s ] → [ √ s ] be a hash function. We define ˜ S ∈ R s × √ s to be

That is, each column i of ˜ S is the sum of the standard basis vectors e j such that j is mapped by h to i . In other words, the locations of the non-zero entries in column i correspond to the preimage of i under h . We then multiply each non-zero entry of ˜ S independently at random by ± 1 . Next, we replace the k th row in ˜ S by multiplying all non-zero ± 1 entries at index i with the binary representation of i to get a matrix S ( m ) ∈ R s × ( √ s × log s ) . That is, for a non-zero entry at index i in row k , we replace the i th entry with ± 1 · B i . Note here that each column still has at most √ s non-zero entries. Finally, we stack 243 · log s -many copies of S ( m ) as heads so that each copy produces independent measurements x S ( m ) . Here, we want to emphasize that each such copy uses fresh randomness for multiplying the non-zero entry of ˜ S independently at random by ± 1 .

Now, we will show that the average of the measurements with matrices S ( m ) ∈ R s × √ s log s can locate the heavy-hitter in x , where x is the vector of inner products from (C.12). For this purpose, we first specify the algorithm for decoding the heavy-hitter.

Algorithm 2 Decoder

Require: The vector y such that y = x S .

2: Take the average y ← 1 243 · log s ∑ k y ( m ) .

1: Split y into 243 · log s blocks y ( m ) ∈ R √ s log s , each of which is a result of multiplying x by S ( m ) , m ∈ [243 · log n ] .

3: b ← I r ( | y | ) ∈ R √ s log s , cf. (C.14).

4: Retrieve b by isolating the binary representation of the position of the heavy-hitter in x .

Here, we define the function I r : R √ s log s → Z √ s log s 2 to [ x S ] m that rounds each entry of its input to the nearest integer:

That is, in both cases, we retrieve the row in S ( m ) that corresponds to the heavy-hitter in x . Since the rows in S ( m ) are distinct, we can also infer the position of the heavy-hitter in x with probability 1.

We now show that y = x S , which consists of 243 · log n many independent copies of y ( m ) = x ˆ S ( m ) . Instead, notice that we can analyze ˜ y = x ˜ S since each ˆ y is the replacement of non-zero entries of ˜ S with the binary representation of their indices times y . We will drop the superscript for now to avoid cluttering the notation. We can then make the following claim:

̸

For the above claim, we note that the first part follows from the latter as it suffices to show that all the non-zero heavy-hitters contribute O ( ϵ 4 √ s ) to the sum ˜ y i = ⟨ x, ˜ S [: , i ] ⟩ . Since each column in ˜ S only interacts with √ s sized sub-vector of x , each ˜ y j for non-heavy hitters can be expressed as

√

where S j ∈ R s contains the non-zero entries of ˜ S j and x is obtained by extracting the entries with corresponding indices from x . Here, we have ∥ x ∥ 2 ≤ ϵ 4 √ s since each entry associated with the nonheavy hitter is bounded as x i ≤ ϵ , and thus, ∥ x ∥ 2 = √∑ i x 2 i ≤ √∑ i ϵ 2 = √ √ s · ϵ 2 = ϵ 4 √ s . Consequently, as ˜ S j is independently random ± 1 , we then must have

with constant probability for j = i . To see this, note that

2

̸

= ∥ x ∥ 2 , where the last equality follows since E [ S jk · S jℓ ] = δ k,ℓ by the distribution on entries of S j . Now, we use Jensen's inequality [57] to get the following bound on the expectation of ∣ ∣ ⟨ x, S j ⟩ ∣ ∣ .

We then use the expectation above to bound the relevant probability as follows:

̸

Now, define the average ˜ y j := 1 243 · log s ∑ m ˜ y ( m ) so that y j (line 3 in Decoder ) is the corresponding replacement of the non-zero entries with the binary representation of their indices. We now claim that this average ˜ y j ≤ 4 / 9 < 1 / 2 with high probability for j = i . To this end, we employ the multiplicative where we apply Markov's inequality [57] in (C.18). That is, we have shown that ˜ y j is bounded by 1 / 3 with constant probability for j = i , and ˜ y i is thus bounded by 1 ± 1 3 . Note here that each of the m -copies ˜ y ( m ) i will have identical guarantees.

̸

Chernoff bound [57] on the independent random variables { ˜ y ( h ) j } h [0 , 1] with E [ ∑ m ˜ y ( m ) j ] ≤ 81 · log s to get

̸

̸

̸

Therefore, we have shown that the average ˜ y j is less than 1 / 2 with probability at least 1 -1 s 3 for j = i . Consequently, we will have ˜ y i bounded by 1 ± 1 2 . Using the union bound over each j = i and the log s bits in the binary representation of j , we can then show that y j + m < 1 / 2 for each j = i, m ∈ [0 , log s ] with probability 1 -log s s 2 ≫ 1 -1 s .

D Experimental Details

D.1 Pre-training

To verify the effect of introducing heads to Hyena as described in Section 4, we train a series of models on THE PILE [11]. All MultiHyena models are set to 8 heads, and otherwise use the same hyperparameters of Hyena models of equivalent size. We set weight decay of Hyena filter parameters to 0 , and lower the frequency of sine activations in the implicit MLP to 4 . We follow the setup of [2], and first train models for 5 , 10 and 15 billion tokens, adjusting the learning rate scheduler accordingly. Then, we train for 300 billion tokens. The results are reported in Tables 5.1 and 5.1.

D.2 Distillation Analysis

Distilling pre-trained long convolution sequence models ( LCSM ) with LaughingHyena can introduce errors on the convolution filter, which then propagate to the outputs.

Setup We perform a series of extensive experiments on all variants of LCSM , including pre-trained H3 models of sizes 125 million, 355 million, 1 . 3 billion and 2 . 7 billion parameters; Hyena of size 153 million parameters, and MultiHyena of size 153 million parameters. For H3 models, we report approximation errors on both shift as well as diagonal SSMs (reported as IIR and FIR). Each point corresponds to distillation carried out at a particular order, using LaughingHyena modal interpolation. To optimize the parameters of the modal form, we use gradient-based optimization and minimize the ℓ 2 discrepancy between filters in time domain. In particular, we use the ADAMW [58] optimizer with learning rate 3 · 10 -4 , and a cosine annealing decay schedule down to 10 -6 after 30 thousand iterations. Each individual filter of every layer is distilled in the same way.

Discussion The errors are shown in Figures D.1, D.1, D.2, D.3, D.4 and D.5. We observe H3 filters to be easier to distill into recurrences with small state without introducing significant errors, whereas Hyena variants learn filters with larger effective dimensions. This provides further evidence that training with implicit convolutions may yield in general more expressive filters.

Distillation Error of H3 125M

IIR

2

1

2

3

0

.

00005

0

.

00010

0

.

00015

0

.

00020

Error

IIR

/lscript

1

Error

2

1

2

3

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

×

10

-

6

2

2

1

2

3

0

.

0025

0

.

0050

0

.

0075

0

.

0100

0

.

0125

IIR

/lscript

∞

Error

2

0

2

2

2

4

0

.

05

0

.

10

0

.

15

0

.

20

IIR Hankel Singular Values

2

1

2

3

2

5

Order

0

.

0005

0

.

0010

0

.

0015

0

.

0020

0

.

0025

Error

FIR

/lscript

1

Error

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

×

10

-

5

FIR

/lscript

2

Error

2

1

2

3

2

5

Order

0

.

004

0

.

006

0

.

008

0

.

010

FIR

/lscript

∞

Error

2

0

2

3

Singular Value Index

0

.

1

0

.

2

0

.

3

FIR Hankel Singular Values

/lscript

Error

D.2.1 Pretrained Filters: Effective Dimension

Visualizations We qualitatively investigate LCSM filters at initialization and after pretraining. This visual inspection (Figures D.6, D.7 and D.8) complements the distillation error analysis of Section D.2.

Distribution of Hankel singular values We further compute the distribution of Hankel singular values of each long convolution filter in different models. The decay in the spectrum quantifies how easy it is to find a compact modal form with LaughingHyena , and serves as a proxy measure of effective dimension of the convolution. The results are shown in Figures D.9 and D.10.

Distillation Error of H3 355M

2

1

2

3

0

.

00004

0

.

00006

0

.

00008

0

.

00010

Error

IIR

/lscript

1

Error

2

1

2

3

0

2

4

6

×

10

-

7

IIR

/lscript

2

Error

2

1

2

3

0

.

002

0

.

004

0

.

006

0

.

008

0

.

010

IIR

/lscript

∞

Error

2

0

2

2

2

4

0

.

05

0

.

10

IIR Hankel Singular Values

2

1

2

3

2

5

Order

0

.

0010

0

.

0015

0

.

0020

0

.

0025

Error

FIR

/lscript

1

Error

2

1

2

3

2

5

Order

0

.

5

1

.

0

1

.

5

2

.

0

×

10

-

5

FIR

/lscript

2

Error

2

1

2

3

2

5

Order

0

.

004

0

.

006

0

.

008

0

.

010

0

.

012

FIR

/lscript

∞

Error

2

0

2

3

Singular Value Index

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

FIR Hankel Singular Values

Distillation Error of H3 1 . 3B

2

1

2

3

0

.

00006

0

.

00008

0

.

00010

0

.

00012

0

.

00014

Error

IIR

/lscript

1

Error

2

1

2

3

2

4

6

8

×

10

-

7

IIR

/lscript

2

Error

2

1

2

3

0

.

004

0

.

006

0

.

008

0

.

010

0

.

012

IIR

/lscript

∞

Error

2

0

2

2

2

4

0

.

00

0

.

05

0

.

10

IIR Hankel Singular Values

2

1

2

3

2

5

Order

0

.

0010

0

.

0015

0

.

0020

Error

FIR

/lscript

1

Error

2

1

2

3

2

5

Order

0

.

5

1

.

0

×

10

-

5

FIR

/lscript

2

Error

2

1

2

3

2

5

Order

0

.

004

0

.

006

0

.

008

FIR

/lscript

∞

Error

2

0

2

3

Singular Value Index

0

.

05

0

.

10

FIR Hankel Singular Values

D.3 Downstream Evaluation

We benchmark the downstream performance of MultiHyena and distilled MultiHyena on standard language modeling tasks from the LM-Eval-Harness [42] and HELM [41] suites. As a reference baseline, we evaluate Pythia [44] 160 M.

Our objective is to quantify the absolute performance of MultiHyena and the downstream impact of distillation. We use the same procedure outlined in Section D.2 to distill MultiHyena .

D.4 Benchmarking

To demonstrate the superior performance of Laughing Hyena for autoregressive generation, we conduct a series of experiments to benchmark its latency, throughput, and memory usage for autoregressive generation with initial prompt length T and number of generated tokens K . For each experiment, we compare the performance of Laughing Hyena against a Transformer, a hybrid H3-attention model with 2 attention layers and a Hyena model. The latter two have been shown to match or achieve lower perplexity than Transformers on standard datasets (WIKITEXT103 and THE PILE). All experiments are carried out on a NVIDIA A100 with 80GB in float16 precision. Missing measurements for any model indicate Out of Memory (OOM) errors while doing autoregressive inference for that particular model.

Distillation Error of H3 2 . 7B

2

1

2

3

0

.

00010

0

.

00015

0

.

00020

0

.

00025

Error

IIR

/lscript

1

Error

2

1

2

3

0

.

5

1

.

0

×

10

-

6

IIR

/lscript

2

Error

2

1

2

3

0

.

0050

0

.

0075

0

.

0100

0

.

0125

IIR

/lscript

∞

Error

2

0

2

2

2

4

0

.

00

0

.

05

0

.

10

IIR Hankel Singular Values

2

1

2

3

2

5

Order

0

.

0010

0

.

0015

0

.

0020

0

.

0025

0

.

0030

Error

FIR

/lscript

1

Error

2

1

2

3

2

5

Order

1

2

3

×

10

-

5

FIR

/lscript

2

Error

2

1

2

3

2

5

Order

0

.

006

0

.

008

0

.

010

0

.

012

FIR

/lscript

∞

Error

2

0

2

3

Singular Value Index

0

.

05

0

.

10

FIR Hankel Singular Values

Distillation Error of Hyena 155M ( 150B tokens )

2

1

2

4

2

7

Order

0

.

00

0

.

01

Error

/lscript

1

Error

2

1

2

4

2

7

Order

0

.

000

0

.

005

0

.

010

0

.

015

Error

/lscript

2

Error

2

1

2

4

2

7

Order

0

1

2

Error

/lscript

∞

Error

2

1

2

4

2

7

5

10

Hankel Singular Values

2

1

2

5

Order

0

.

0

0

.

1

0

.

2

Error

/lscript

1

Error

2

1

2

5

Order

0

1

2

3

Error

/lscript

2

Error

2

1

2

5

Order

0

10

20

Error

/lscript

∞

Error

2

1

2

5

50

100

Hankel Singular Values

Distillation Error of MultiHyena 155M

(

300B tokens

)

Peak throughput We first evaluate the throughput (number of tokens generated per second) across different batch sizes, using a typical generation workload consisting of a prompt of length 512 and generating 256 tokens. Figure 1.1 measures peak throughput of different models. Since Laughing Hyena does not require caching intermediate kv-projections during generation, reduced memory requirements at a fixed model size allow it to process larger batch sizes.

Prompt length Autoregressive generation in Laughing Hyena is achieved through a two-step process: an initial prefill step that uses the length -T prompt to initialize the state x T and that generates all K tokens. In Figure 5.3 we demonstrate how the prefill step scales for different prompt lengths, keeping batch sizes fixed at 64 . Since prefilling in Laughing Hyena is carried out efficiently via convolutions (as described in Section 3.4), throughput scales more favorably than Transformers. Other models capable of prefilling via convolutions also achieve higher throughputs than Transformers but are ultimately slower than Laughing Hyena during the generation phase.

State throughput We measure the impact of SSM state dimension on the throughput of Laughing Hyena . Keeping batch sizes fixed reveals minimal impact for all dimensions smaller than

t

FIR Filters of H3 125M

0

50

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 0

Initialization

Pre

-

trained

0

50

-

1

0

1

Layer 1

0

50

-

1

0

1

Layer 2

0

50

-

1

0

1

Layer 3

0

50

-

1

0

1

Layer 4

0

50

-

1

0

1

Layer 5

0

50

t

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Attention Layer

0

50

t

-

1

0

1

Layer 7

0

50

t

-

1

0

1

Layer 8

0

50

t

-

1

0

1

Layer 9

0

50

t

-

1

0

1

Layer 10

0

50

t

-

1

0

1

Layer 11

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 0

Initialization

Pre

-

trained

0

2000

-

1

0

1

Layer 1

0

2000

-

1

0

1

Layer 2

0

2000

-

1

0

1

Layer 3

0

2000

-

1

0

1

Layer 4

0

2000

-

1

0

1

Layer 5

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Attention Layer

0

2000

-

1

0

1

Layer 7

0

2000

-

1

0

1

Layer 8

0

2000

-

1

0

1

Layer 9

0

2000

-

1

0

1

Layer 10

0

2000

-

1

0

1

Layer 11

IIR Filters of H3 125M

t

t

t

t

t

100 , which are sufficient to distill all models discussed in this work. All other measurements provided in this Section are carried out with a standard order 16 . We note that it may be possible to further increase peak throughput by leveraging reduced memory footprints achieved by extremely small SSMs.

Latency over sequence length We benchmark the time taken to generate a variable number of tokens, starting from a prompt of length 512 tokens at batch size 1 (Figure D.11). Laughing Hyena tracks highly optimized Transformers. We note that Laughing Hyena is asymptotically more efficient than Transformers; however, this regime is bottlenecked by hardware-specific implementation details and optimizations. We expect optimized, platform-specific implementations of Laughing Hyena to outperform Transformers even at batch size 1. When the prompt is long, the prefilling step becomes the bottleneck, and all convolutional models outperform Transformers.

Parameter scaling To better understand how the performance of Laughing Hyena scales, we benchmark its latency, throughput, and peak memory utilization for autoregressive generation and 125M, 355M, 1.3B, 2.7B and 6.7B parameters. We compare the performance to that of Transformers, HybridH3, and Hyena at the same number of parameters and report the results in Figure D.11. For the latency measurement, we use a batch size of 1 and benchmark the time taken to generate 128 tokens, starting from a prompt of length 512 tokens. For throughput and peak memory scaling against the number of parameters, we use a batch size of 64 and measure the throughput for generating 256 tokens starting with a prompt of length 512.

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 0

Initialization

Pre

-

trained

0

2000

-

1

0

1

Layer 1

0

2000

-

1

0

1

Layer 2

0

2000

-

1

0

1

Layer 3

0

2000

-

1

0

1

Layer 4

0

2000

-

1

0

1

Layer 5

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 6

0

2000

-

1

0

1

Layer 7

0

2000

-

1

0

1

Layer 8

0

2000

-

1

0

1

Layer 9

0

2000

-

1

0

1

Layer 10

0

2000

-

1

0

1

Layer 11

0

2000

t

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 12

0

2000

t

-

1

0

1

Layer 13

0

2000

t

-

1

0

1

Layer 14

0

2000

t

-

1

0

1

Layer 15

0

2000

t

-

1

0

1

Layer 16

0

2000

t

-

1

0

1

Layer 17

Filters of MultiHyena 155M

(

300B tokens

)

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 0

Initialization

Pre

-

trained

0

2000

-

1

0

1

Layer 1

0

2000

-

1

0

1

Layer 2

0

2000

-

1

0

1

Layer 3

0

2000

-

1

0

1

Layer 4

0

2000

-

1

0

1

Layer 5

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 6

0

2000

-

1

0

1

Layer 7

0

2000

-

1

0

1

Layer 8

0

2000

-

1

0

1

Layer 9

0

2000

-

1

0

1

Layer 10

0

2000

-

1

0

1

Layer 11

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 12

0

2000

-

1

0

1

Layer 13

0

2000

-

1

0

1

Layer 14

0

2000

-

1

0

1

Layer 15

0

2000

-

1

0

1

Layer 16

0

2000

-

1

0

1

Layer 17

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 18

0

2000

-

1

0

1

Layer 19

0

2000

-

1

0

1

Layer 20

0

2000

-

1

0

1

Layer 21

0

2000

-

1

0

1

Layer 22

0

2000

-

1

0

1

Layer 23

0

2000

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 24

0

2000

-

1

0

1

Layer 25

0

2000

-

1

0

1

Layer 26

0

2000

-

1

0

1

Layer 27

0

2000

-

1

0

1

Layer 28

0

2000

-

1

0

1

Layer 29

0

2000

t

-

1

0

1

Normalized Filters

h

t

/

‖

h

‖

∞

Layer 30

0

2000

t

-

1

0

1

Layer 31

0

2000

t

-

1

0

1

Layer 32

0

2000

t

-

1

0

1

Layer 33

0

2000

t

-

1

0

1

Layer 34

0

2000

t

-

1

0

1

Layer 35

Filters of Hyena 355M

(

200B tokens

)

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Hankel Singular Values

Layer 0

MultiHyena

Hyena

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 1

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 2

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 3

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 4

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 5

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Hankel Singular Values

Layer 6

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 7

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 8

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 9

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 10

2

0

2

4

2

8

10

-

4

10

-

2

10

0

10

2

Layer 11

2

0

2

4

2

8

n

10

-

4

10

-

2

10

0

10

2

Hankel Singular Values

Layer 12

2

0

2

4

2

8

n

10

-

4

10

-

2

10

0

10

2

Layer 13

2

0

2

4

2

8

n

10

-

4

10

-

2

10

0

10

2

Layer 14

2

0

2

4

2

8

n

10

-

4

10

-

2

10

0

10

2

Layer 15

2

0

2

4

2

8

n

10

-

4

10

-

2

10

0

10

2

Layer 16

2

0

2

4

2

8

n

10

-

4

10

-

2

10

0

10

2

Layer 17

Hankel Singular Values of Hyena and MultiHyena Filters

Hankel Singular Values of H3 ( 125M ) FIR Filters

2

0

2

4

10

-

5

10

-

2

Hankel Singular Values

Layer 0

Initialization

Pre

-

trained

2

0

2

4

10

-

5

10

-

2

Layer 1

2

0

2

4

10

-

5

10

-

2

Layer 2

2

0

2

4

10

-

5

10

-

2

Layer 3

2

0

2

4

10

-

5

10

-

2

Layer 4

2

0

2

4

10

-

5

10

-

2

Layer 5

0

.

0

0

.

5

1

.

0

n

-

1

0

1

Hankel Singular Values

Attention Layer

2

0

2

4

n

10

-

5

10

-

2

Layer 7

2

0

2

4

n

10

-

5

10

-

2

Layer 8

2

0

2

4

n

10

-

5

10

-

2

Layer 9

2

0

2

4

n

10

-

5

10

-

2

Layer 10

2

0

2

4

n

10

-

5

10

-

2

Layer 11

2

0

2

4

10

-

5

10

-

2

Hankel Singular Values

Layer 0

Initialization

Pre

-

trained

2

0

2

4

10

-

5

10

-

2

Layer 1

2

0

2

4

10

-

5

10

-

2

Layer 2

2

0

2

4

10

-

5

10

-

2

Layer 3

2

0

2

4

10

-

5

10

-

2

Layer 4

2

0

2

4

10

-

5

10

-

2

Layer 5

0

.

0

0

.

5

1

.

0

n

-

1

0

1

Hankel Singular Values

Attention Layer

2

0

2

4

n

10

-

5

10

-

2

Layer 7

2

0

2

4

n

10

-

5

10

-

2

Layer 8

2

0

2

4

n

10

-

5

10

-

2

Layer 9

2

0

2

4

n

10

-

5

10

-

2

Layer 10

2

0

2

4

n

10

-

5

10

-

2

Layer 11

Hankel Singular Values of H3

(

125M

)

IIR Filters

0

1

,

000

2

,

000

0

20

40

Output sequence length

Full gener. latency (s)

0

2

4

6

0

20

40

60

Number of Parameters

×

10

9

Peak Memory [GBs]

0

1

2

0

2

4

6

Number of Parameters

×

10

9

Gener. throughput [tok/s]

×

10

3

0

2

4

6

2

4

Number of Parameters

×

10

9

Full gener. latency (s)

Laughing Hyena

Hyena

Transformer

E Additional Experiments

E.1 Associative Recall with MultiHyena

We follow the setup of [2] and train 2-layer Hyena and MultiHyena (with 8 heads) to solve associative recall via a standard next-token prediction objective. We focus on the sequence length 64 k, high vocabulary size setting, and push vocabulary sizes past the maximum values considered in [2]. At vocabulary size 60 , a difference between MultiHyena and Hyena can be observed (Table E.1), as experimental support for Theorem 4.1.

E.2 Analysis of Hankel Singular Values of Pretrained Large Convolution Sequence Models

E.3 Model Order Reduction of H3

The H3 model is constructed with a combination of diagonal SSMs and shift SSMs. There exists various classical model order reduction techniques for these different types of SSMs. The following sections aim to present the formulation and effectiveness of two classical approaches on obtaining the compressed representation of a H3 model. More specifically, we study modal truncation and balanced truncation for compressing diagonal SSMs and shift SSMs respectively.

E.3.1 Modal Truncation

A discrete diagonal SSM ( A = diag ( λ 1 , . . . , λ d ) , B ∈ C d × 1 , and C ∈ C 1 × d ) can be directly converted into a residue-pole transfer function as follows:

where residue r i = B i C i . Modal truncation aims to compress such a transfer function by essentially reducing the summation over d to n < d , of the n most influential modes. The influence from each node can be isolated by expressing it using the h ∞ norm of the system as follows:

Each mode i can be ranked using the bound formulated above. Subsequently, the d -n lowest modes could be discarded to form a reduced order model. Figure E.1 illustrates the monotonically decreasing l ∞ error with the increase in system order. However, this model reduction approach is only suitable for diagonalizable SSMs.

Modal Truncation Model Reduction Error of H3

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 0

Mean

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 1

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 2

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 3

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 4

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 5

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 6

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 7

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 8

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 9

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 10

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 11

E.3.2 Balanced Truncation

A balanced SSM realization is one in which the observability P and controllability Q gramians are equal and diagonal. Such a realization can be formulated with the following Lyapunov equations [24]:

where P = Q = Σ = diag ( σ 1 , . . . , σ d ) and σ i ≥ σ i +1 .

Results from [59] shows that the n -order model reduction error is bounded by:

Therefore, an n -order partition of the full balanced realization can be chosen, such that the discarded orders are the d -n lowest contributor to the error. The steps taken by [24], computes the n -order partition of the balanced realization as follows:

Form a Hankel matrix S d from the impulse response h 1: n of the shift SSM.

Obtain the eigenvector matrix V ∈ C d × d , and the eigenvalues λ = σ 2 via the eigendecomposition of S d .

Choose the truncated model's order n < d , based on the bound in Equation E.4.

Compute the state-space matrices as follows:

This model reduction technique was applied to a trained 125 million parameter H3 , MultiHyena , and Hyena models as shown in Figures E.2, E.3, and E.4 respectively. It could be noted that the all models encountered an undesirable non-monotonic error reduction with the increase in order. Moreover, order reduction configurations such as the one in Figure E.3 layer 15 display signs of numerical instability.

Balanced Truncation Model Reduction Error of H3

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 0

Mean

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 1

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 2

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 3

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 4

2

1

2

3

2

5

0

.

0

0

.

5

1

.

0

Layer 5

2

1

2

3

2

5

Order

0

.

00

0

.

02

0

.

04

Normalized

/lscript

∞

Error

Layer 6

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 7

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 8

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 9

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 10

2

1

2

3

2

5

Order

0

.

0

0

.

5

1

.

0

Layer 11

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 0

Mean

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 1

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 2

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 3

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 4

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 5

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 6

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 7

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 8

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 9

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 10

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 11

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 12

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 13

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 14

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 15

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 16

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 17

Balanced Truncation Model Reduction Error of MultiHyena

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 0

Mean

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 1

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 2

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 3

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 4

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 5

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 6

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 7

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 8

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 9

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 10

2

1

2

4

2

7

0

.

0

0

.

5

1

.

0

Layer 11

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Normalized

/lscript

∞

Error

Layer 12

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 13

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 14

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 15

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 16

2

1

2

4

2

7

Order

0

.

0

0

.

5

1

.

0

Layer 17

Balanced Truncation Model Reduction Error of Hyena

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis article addresses two issues:\n\n1. The low efficiency of the LonvConv-based model during inference.\n2. Whether it is advantageous to perform independent long convolutions on each channel or reduce the total number of filters without loss in quality.\n\nTo tackle problem (1), the authors propose distilling the LonvConv into a Diagonal State space model and train it using the $\\ell_2$ loss function.\n\nFor problem (2), the authors suggest sharing long convolution coefficients across multiple channels, resulting in the MultiHyena structure.\n\nThe effectiveness of the proposed methods is validated on multiple datasets.\n\nSTRENGTHS:\nDistilling LonvConv into a Diagonal State space model is indeed a novel and meaningful approach. The conclusion of sharing long convolution coefficients across multiple channels is also innovative.\n\nWEAKNESSES:\nI think the main issue with this article lies in the focus of the writing. From Equation 3.4, it is clear that the ultimate goal is to represent the LonvConv coefficients using a Diagonal State Space Model. However, a significant portion of the article is spent describing unrelated aspects. The most crucial part, determining the hidden dimension of the State Space Model, is merely brushed over, even though it is the key factor that affects the quality of the distillation and the efficiency of the final inference. On the other hand, the motivation behind the design of MultiHyena should be addressed in the main text since it is of utmost importance.\n\nQUESTIONS:\n1. The method for determining the hidden dimension of the State Space Model needs to be explained in more detail. I referred to Appendix E.3.2, and I'm wondering if the core idea is to perform an SVD decomposition and then select the dimension for dimensionality reduction based on the eigenvalues?\n\n2. The solution to Equation 3.4 requires a more detailed algorithm description or pseudocode to help readers follow along. I attempted to replicate it following Appendix B.1, but the results were not quite good. Could you provide the training configurations and an estimate of the training time?\n\n3. The motivation behind the design of MultiHyena should be included in the main text, and ablation studies (sharing coefficients vs not sharing coefficients) should be conducted to validate the design's rationale. The experiments should compare the effects and speeds.\n\n4. Does MultiHyena utilize the Local conv + Global conv structure of Hyena? If so, how many layers are used? This should be mentioned in the experiments.\n\n5. Is the Algorithm 1 on page 28 is the implementation of Figure 4.1?\n\n6. Regarding the implementation of MultiHyena, in Algorithm 1 on page 28, for $z^m_t \\in \\mathbb R^{L\\times N\\times N}$, what does the subscript $t$ represent? On the other hand, is the shape of $T_h$ $L$ (all features share one set of convolution coefficients) or $L\\times N\\times N$ (each feature has independent convolution coefficients)?\n\n7. Continuing with the implementation of MultiHyena, in Algorithm 1 on page 28, is the shape of $T_hz^m_t$ ${L\\times N\\times N}$? If so, does $T_h(z^m_t)  q_t^m$ mean performing matrix multiplication between $[T_h(z^m_t)]_i \\in \\mathbb R^{N\\times N}$ (for $i=0,\\ldots,L-1$) and $q_t^m\\in \\mathbb R^{N}$, resulting in an output of shape $\\mathbb R^{N}$? If not, what is the computation like?\n\n8. The algorithm's output is $\\bar{y} \\leftarrow\\left(\\sum_m\\right) y^m / M\\in \\mathbb R^{L\\times N}$, while the input has a shape of $L\\times D$. Is this inconsistency problematic, or did the algorithm omit something?\n\n9. Algorithm 1 has several ambiguities. It is suggested to reorganize it for better clarity.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes distilling convolutional models for autoregressive sequence generation into recurrent (state-space) models. A key limitation of recently proposed convolutional models is that they use convolutional filters that extend potentially infinitely into the past, and the same techniques that yield efficient training do not transfer over to token-by-token generation. This paper proposes a post-training and data-free distillation step that replaces such convolutional models with a recurrent approach that uses constant time and memory at each step of generation during inference.\n\nSTRENGTHS:\nThe key strength of the paper is that proposes a novel and theoretically grounded distillation method, which achieves good approximation bounds without being tied to specific data or involving what amounts to an additional round of training.\n\nAlso, convolutional and state space models as a whole are an underexplored area in recent work when compared to Transformers and attention, and this paper puts forward a novel method of linking the two, both of which contribute to the originality of the paper.\n\nWEAKNESSES:\nMy biggest reservation based on my understanding of the paper is that I didn't get a qualitative sense of what might be lost as part of the distillation process. Is there some intuition of what is the worst-case qualitative behavior of a convolutional filter that can't be tightly approximated with LaughingHyena?\n\nThis question comes to mind because in the world of Transformers, it is no secret that most of the computational power spent on quadractic attention goes to waste. Just a small subset of the efficiency work there includes pruning entire heads, limiting each head to a head-specific attention context window, or even doing sliding-window attention with a fixed context length for the entire model. More in line with the present paper, work like Performer has developed an approximation for converting attention-based models into recurrent models -- at a cost. Notwithstanding the theory of the tightness of that last approximation, and equivalent performance of many to the Transformer that is demonstrated in some of the papers introducing these methods, there inevitably arises some situation where none of the approximations match the Transformer in quality. I worry that the present approach might fall into a similar pattern. A recurring theme in this area is that any method that sacrifices the ability to have long context, or to perform associative recall, is suspect. That's why when it comes to these approximations of convolutions, it would helpful to know whether the approximation is effectively some form of context-truncation in disguise, and if not what the qualitative cost is.\n\nQUESTIONS:\nHow well does the LaughingHyena architecture perform on the associative recall task, especially as compared to Hyena (or MultiHyena)? Is the point beyond which the models fail to perform the task (in terms of sequence length or vocabulary size) different between the two?\n\nFor LM-Eval-Harness and HELM, have you tried a baseline of taking the impulse response from Hyena/MultiHyena and truncating it to a finite impulse response? A sliding window seems like one of the simplest approximations to try in the world of convolutions, and it would be helpful to know if maybe some defect of the tasks or the base model results in nothing more being required. This would, in fact, be a useful baseline to have in the paper.\n\nLIMITATIONS:\nThe paper could be improved with a little bit of additional discussion regarding limitations of the distillation/approximation.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper proposes LaughingHyena - an improvement to the Hyena model that can perform long-convolutions instead of attentions in transformers to avoid the quadratic scaling issues. One of the issues with the convolution sequence models is that they incur significant cost due to autoregressive inference. To avoid this, this paper seeks to come up with a techinique to have constant memory recurrent inference to increase generation thoroughput. This is achieved using the use of compact linear SSM and weight tying the filters across heads in Hyena architecture. The resulting performance improvements are impressive - 1.5x throughput improvement compared to Hyena. The model also achieves SOTA in the PILE dataset.

STRENGTHS:
- The perplexity on small-scale models on Table 1 and Table 2 outperform GPT, Hyena and establishes a new SOTA.
- The peak memory usage is also constant for different sequence lengths in Table 5.4

WEAKNESSES:
- The writing is a bit hard to follow.
- The performance of the model is still lacking compared to full transformer baseline such as Pythia in Table 5.3. Can the authors comment on this? Any idea on how much the performance degradation will be on very large scale models and datasets?

QUESTIONS:
- What assumptions do you use for the state-space model in Eq 3.1 to yield a good distillation results (d<<L)?

LIMITATIONS:
I think the writing can be improved to provide a simple explanation of the method for readers who don't have a strong understanding of state-space models. Else, the text is hard to follow.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09

comments: SUMMARY:
This paper proposes an approach that enables constant-memory, recurrent inference for long convolution architectures. They introduce LaughingHyena, a distilation approach that consists of extracting compact linear SSMs from each convolution layer. Combined with weight-tying, it results in state-of-the-art performance and efficiency (i.e. throughput) without any drop in quality.

STRENGTHS:
The paper is well structured and written.

The approach seems sound, reasonable and is performant

WEAKNESSES:
The models used in most experiments are small.

The helm evaluation is not very convincing.

QUESTIONS:
Is it possible to benchmark against more recent open source models such as Llamma?

LIMITATIONS:
Broader Impacts section is missing.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-08-04