SlowGuess/NeurIPS_2023_OCR huggingface dataset:
LAUGHING HYENA DISTILLERY Extracting Compact Recurrences From Convolutions
Stefano Massaroli*,1, Michael Poli*,2, Daniel Y. Fu*,2, Hermann Kumbong², Rom N. Parnichkun³, Aman Timalsina⁴, David W. Romero⁵, Quinn McIntyre², Beidi Chen⁶, Atri Rudra⁷, Ce Zhang⁸, Christopher Re²;†, Stefano Ermon²;†, Yoshua Bengio¹;†
$^{*}$ Equal contribution. † Equal senior authorship.  $^{1}$ Mila and Université de Montréal.  $^{2}$ Stanford University.  $^{3}$ The University of Tokyo.  $^{4}$ Purdue University.  $^{5}$ Vrije Universiteit Amsterdam.  $^{6}$ Carnegie Mellon University and Meta AI (FAIR).  $^{7}$ University of Buffalo, SUNY.  $^{8}$ University of Chicago and Together Computer.
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 pretraining quality and reduce the number of filters to be distilled. The resulting model achieves  $10\times$  higher throughput than Transformers and  $1.5\times$  higher than Hyena at 1.3B parameters, without any loss in quality after distillation.
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  $\mathcal{O}(dK)$  time and with constant  $\mathcal{O}(d)$  memory - improving over the  $\mathcal{O}(K^2)$  time and  $\mathcal{O}(K)$  memory usage of kv-cached Transformers and naively executed long convolutions. At model sizes above one billion parameters, LaughingHyena achieves  $10\times$  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\times$  less memory than with a Transformer. At smaller batch sizes, latency of LaughingHyena is also competitive with Transformers, reaching  $\geq 2\times$  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].
Figure 1.1: Throughput (in generated tokens) of Transformers, H3 and Hyena models. LaughingHyena is a recurrent model distilled from a pre-trained Hyena. Workload involves generating 256 tokens given a prompt of length 512.
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$ :
$$
(h * u) _ {t} = \sum_ {j = 0} ^ {t} h _ {t - j} u _ {j}. \tag {2.1}
$$
Generally,  $u_{t} \in \mathbb{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 \mapsto 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  $\gamma_{\theta}: t \mapsto h_t$ .
State-space realization One option is to select  $\gamma_{\theta}$  as the impulse response function of a discrete linear time-invariant system,
$$
\begin{array}{r} x _ {t + 1} = \mathrm {A} x _ {t} + \mathrm {B} u _ {t} \\ y _ {t} = \mathrm {C} x _ {t} + h _ {0} u _ {t}, \quad t \mapsto h _ {t} = \left\{ \begin{array}{l l} h _ {0} & t = 0 \\ \mathrm {C A} ^ {t - 1} \mathrm {B} & t > 0 \end{array} \right. \end{array} \tag {2.2}
$$
with state  $x_{t} \in \mathbb{R}^{d}$ , input  $u_{t} \in \mathbb{R}$ , and output  $y_{t} \in \mathbb{R}$ . The matrices  $\mathsf{A} \in \mathbb{R}^{d \times d}$ ,  $\mathsf{B} \in \mathbb{R}^{d \times 1}$ ,  $\mathsf{C} \in \mathbb{R}^{1 \times d}$ , and  $h_0 \in \mathbb{R}$  are the learnable parameters of the model while the initial state  $x_0$  is usually set to zero such that  $u \mapsto 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  $\gamma_{\theta}(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.
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 data-controlled gating to replace attention with sub-quadratic scaling in sequence length<sup>1</sup>. We use the deep learning convention of naming different projections as query  $q$ , key  $k$  and value  $v$ . Let  $\mathsf{M}_q$  and
$\mathsf{M}_k$  be the  $L$ -by- $L$  diagonal matrices whose respective main diagonal entries are the respective entries of length- $L$  sequences  $q$  and  $k$ . A  $\mathsf{H}$ -block realizes a surrogate attention matrix with a data-controlled, parameterized decomposition in three terms:
Figure 2.1: H-block.  $\mathsf{T}^{(q)},\mathsf{T}^{(k)},\mathsf{T}^{(v)}$  are short-convolution operators.
$$
(q, k, v) \mapsto \mathrm {H} (q, k) v, \quad \mathrm {H} (q, k) = \mathrm {M} _ {q} \mathrm {T} _ {h} \mathrm {M} _ {k} \tag {2.3}
$$
where  $\mathsf{T}_h\in \mathbb{R}^{L\times L}$  is the Toeplitz matrix constructed from the learnable long convolution filter  $h$ , i.e.,  $\mathsf{T}_h = (h_{i - j})_{i,j = 0}^{L - 1}$ . 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² 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  $\tilde{\mathcal{O}}(L)\coloneqq \mathcal{O}(L\log_2L)$  time (two convolutions³ and two element-wise products),  $y_{t} = q_{t}(h*k\nu)_{t}$ . The overall operator acts on an input  $u$  by constructing a third-order 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 length- $T$  prompt  $u \in \mathbb{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  $\tilde{\mathcal{O}}(T)$  time and obtaining a length- $T$  output  $u \mapsto y_0, \ldots, 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} = \sum_{j=0}^{t} h_{t-j} y_j$  for  $t = T-1, \ldots, 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 length- $T$  prompt has time complexity  $\mathcal{O}(T\log_2T + TK + K^2)$  and requires  $\mathcal{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  $\mathcal{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  $\tilde{\mathcal{O}}(T)$  time and  $\mathcal{O}(T)$  memory to process the prompt<sup>4</sup> and initialize the state  $x_{T-1}$ . Each additional generation step only requires  $\mathcal{O}(d)$  time.
Lemma 2.2. Generating  $K$  tokens with a state-space model (2.2) from a length- $T$  prompt has time complexity  $\mathcal{O}(T\log_2T + dK)$  and requires  $\mathcal{O}(T + d)$  memory.
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.
Transformers Self-attention is certainly less efficient than long convolutions in processing the prompt, coming with a hefty  $\mathcal{O}(T^2)$  time complexity. However, Transformers can achieve a similar
In this work, we consider second-order Hyena blocks [2] to automatically extend our findings to H3 [1].
2The short filters are explicitly parameterized, see [2].
<sup>3</sup>The  $qkv$  short convolutions can be evaluated in batch with a single pass. The second convolution is the one with the long filter  $h$  and performed via Fast Fourier Transform (FFT), hence the  $\tilde{\mathcal{O}}(L)$  complexity.
4In §3.4 we show that multiple pre-filling strategies exist, with different trade-offs in time and memory.
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
$$
y _ {t + 1} = \frac {\sum_ {j = 0} ^ {t + 1} \varphi (q _ {t + 1} k _ {j}) v _ {j}}{\sum_ {i = 0} ^ {t + 1} \varphi (q _ {t + 1} k _ {j})} \quad \text {w h e r e} \varphi : \mathbb {R} \to \mathbb {R} \text {i s u s u a l l y c h o s e n a s} \varphi (x) = e ^ {x}.
$$
Lemma 2.3. Generating  $K$  tokens with self-attention from a length- $T$  prompt has time complexity  $\mathcal{O}(T^2 + TK + K^2)$  and requires  $\mathcal{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, \ldots, h_L$ , find a state-space model (2.2) of dimension  $d \ll 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 An alternative description of the system (2.2) is its transfer function  $H$ , defined as the  $z$ -transform of the impulse response  $H(z) = \sum_{t=0}^{\infty} h_t z^{-t}$  for all  $z \in \mathbb{C}$  where the sum converges. The transfer function is a proper rational function of  $z$
$$
H (z) = h _ {0} + \mathrm {C} (z \mid - \mathrm {A}) ^ {- 1} \mathrm {B} = h _ {0} + \frac {b _ {1} z ^ {- 1} + \cdots + b _ {d} z ^ {- d}}{1 + a _ {1} z ^ {- 1} + \cdots + a _ {d} z ^ {- d}}. \tag {3.1}
$$
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  $\mathbb{C}$ -plane circle of radius  $\rho(\mathsf{A})$ ,  $\mathbb{D}_{\rho(\mathsf{A})} \coloneqq \{z \in \mathbb{C} : |z| > \rho(\mathsf{A})\}$  where  $\rho(\mathsf{A})$  is the spectral radius of  $\mathsf{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  $\mathsf{A}, \mathsf{B}, \mathsf{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  $\tilde{\mathcal{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  $\mathcal{O}(d)$  (Lemma A.7), compared to
Figure 3.1: The LaughingHyena long convolution sequence model distillation blueprint.
the  $\mathcal{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:
1. 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.
2. 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.
3. 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  $\hat{h}_t = \mathsf{CA}^{t-1}\mathsf{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 - \hat{y} \|_r \leq \| h - \hat{h} \|_q \| u \|_p$  for properly chosen norms<sup>5</sup>. 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<sup>6</sup>. We design distillation algorithms which either match filters in time domain minimizing the  $\ell_2$  error  $(\| h \|_2 := [\sum_{t \in \mathbb{Z}} |h_t|^2]^{1/2})$  or match their transfer functions optimally with respect to the  $\mathcal{H}_2$  norm  $(\| H \|_2 := [(1/2\pi)\int_{-\pi}^{\pi}|H(e^{i\omega})|^2\mathrm{d}\omega^{1/2})^7$ . As the distillation is carried out via gradient methods,  $\ell_2$  is a natural candidate.  $\mathcal{H}_2$  error minimization can instead be used to uniformly bound the worst-case discrepancy as  $\| h - \hat{h} \|_{\infty} \leq \| H - \hat{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
$5p, q, r > 0$  should satisfy  $1/q + 1/p = 1/r + 1$ . In the case of infinite sequences defined on the all  $\mathbb{Z}$ , the norms are taken in a  $\ell_p, \ell_q, \ell_r$  sense, respectively. The bound is potentially sharp [27, 28]
For completeness, we also test balanced and modal truncation techniques on a suite of pre-trained H3 and Hyena models in Appendix E.3.
Such norms are always well-defined for finite sequences of interest which are in  $\ell_{\infty}$ .
appealing invariance properties, we can parametrize a proper rational function  $\hat{H}(z)$  (3.1) and fit it to the (truncated) transfer function of the original filter  $H_{L}(z) \coloneqq \sum_{t=0}^{L} 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  $\lambda_{n}$  in the denominator's polynomial,  $\lambda_{n} \in \mathrm{roots}(\mathrm{poly}(a))$ .
Proposition 3.1 ([5]). If  $\mathrm{poly}(a)$  has distinct roots  $\{\lambda_n\in \mathbb{C}\}$ , then the transfer function of the system can be factorized as  $\hat{H} (z) = \sum_{n = 1}^{d}R_{n} / (z - \lambda_{n}),\quad \forall z\in \mathbb{D}_{\rho (\mathbf{A})}$  where  $\{R_{n}\in \mathbb{C}\}$  is the residue associated with the pole  $\lambda_{n}$ .
Computing the inverse transform of the expanded transfer function via, e.g., the Cauchy residue theorem [30], shows that the resulting impulse response  $\hat{h}$  corresponds to a truncated basis of exponentially decaying complex sinusoids
$$
\hat {h} _ {t} = \sum_ {n = 1} ^ {d} R _ {n} \lambda_ {n} ^ {t - 1}, \quad R _ {n}, \lambda_ {n} \in \mathbb {C}, t > 0. \tag {3.2}
$$