SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In practice, this corresponds to the impulse response of state-space model with diagonal matrix  $\mathsf{A} = \mathrm{diag}(\lambda_1,\ldots ,\lambda_d)$  and such that  $\mathtt{B}_i\mathtt{C}_i = R_i$  for all  $i = 1,\dots ,d$ . The distillation problem can be then defined in terms of the  $L$ -point nonlinear least squares interpolation error (squared  $\ell_{2}$ ) between  $h_1,\ldots ,h_L$  and (3.2) evaluated for  $t = 1,\ldots ,L$ :  $\min_{\{\lambda_n,R_n\}}\| \hat{h} -h\| _2^2$ . Note that in case of the target filter  $h$  being real-valued, the objective can be replaced by  $\| \Re [\hat{h}] - h\| _2^2$ .
Figure 3.2: Example of modal interpolation. The approximant is a linear combination of exponentially-decaying complex exponential basis functions with learned decay rate.
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  $\lambda_{n}\coloneqq A_{n}e^{i\theta_{n}}$  and the residues in cartesian form<sup>10</sup>. Note that, with this parametrization we have  $\Re [\hat{h}_t] = \sum_nA_n^{t - 1}[\Re (R_n)\cos (\theta_n(t - 1)) - \Im (R_n)\sin (\theta_n(t - 1))]$ . We can also solve the distillation problem in the  $\mathcal{H}_2$  sense by evaluating  $\hat{h}_t$  and  $h_t$  at  $t = 0,\dots ,L - 1$  and taking their respective
(discrete) Fourier transform before computing the objective. Efficient evaluation of (3.2) is crucial for distillation. In particular we show the following:
Lemma 3.1. Evaluation of  $(\hat{h}_t)_{t=0}^{L-1}$  (3.2) can be done in  $\mathcal{O}(dL)$  time from its modal form and in  $\tilde{\mathcal{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\in \mathbb{N}}$  is the so-called McMillan degree [34]:
$$
d ^ {\star} = \arg \min  _ {d} d: \exists \bar {A} \in \mathbb {C} ^ {d \times d}, B \in \mathbb {C} ^ {d \times 1}, \bar {C} \in \mathbb {C} ^ {\bar {1} \times d} \text {w i t h} h _ {t} = \mathsf {C A} ^ {t - 1} \bar {B}, \forall t > 0 \tag {3.3}
$$
Theorem 3.1 (Ho-Kalman [35, Theorem 2, Corollary]). Let  $\mathsf{S}$  be the (infinite) Hankel matrix constructed with  $h$ , i.e.  $\mathsf{S} := (h_{i+j})_{i,j=1}^{\infty}$ . Then,  $d^{\star} = \mathrm{rank}(\mathsf{S})$ .
As already partially discussed in [6], the truncation introduces a correction term in the approximant transfer function. See Appendix A.4.
${}^{9}$  i.e normalizing denominator polynomial coefficients to constrain roots within the unit circle.
<sup>10</sup>We report additional details on the nuances of the parametrization in Appendix B.1.
A lower bound for  $d^{\star}$  can be estimated from a truncated filter of length  $L$  by constructing the  $L \times L$  principal sub-matrix  $\mathsf{S}_L$  and using the fact that  $\mathrm{rank}(\mathsf{S}) \geq \mathrm{rank}(\mathsf{S}_L)$ . Inspecting how fast the Hankel singular values  $(\sigma_n)_{n=1}^L$  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  $\sigma_{d+1}$  to be sufficiently small[11]. Specifically, we can prove that the last singular value  $\sigma_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 length- $L$  filter,  $\hat{h}$  a distilled filter of order  $d < L$  and let  $S_L, \hat{S}_L$  be the respective Hankel matrices. Then  $\inf_{\hat{S}_L} \| S_L - \hat{S}_L \|_2 = \sigma_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} \in \mathbb{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 length- $T$  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  $\mathcal{O}(dT)$  with a  $\mathcal{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  $\mathcal{O}(d\log_2 T)$  while incurring in an increased memory requirement of  $\mathcal{O}(dT)^{12}$ . A third option is to use a single FFT convolution to obtain  $x_{T}$  in  $\tilde{\mathcal{O}}(T)$  time and  $\mathcal{O}(T)$  memory.
Proposition 3.2.  $x_{T} = (\nu_{T},\dots ,\nu_{T - d})$  where  $\nu_{t} = (g*u)_{t}$  and  $g$  is the filter whose transfer function is  $1 / \mathrm{den}(\hat{H})(z)$  and can be evaluated in  $\tilde{\mathcal{O}} (T)$ .
Note that, the fast pre-filling algorithm established by this result requires evaluating the denominator polynomial of  $\hat{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  $\mathcal{O}(d)$  time and memory:
Proposition 3.3. The filter (3.2) has a state space matrices  $\mathsf{A} = \mathrm{diag}(\lambda_1, \ldots, \lambda_d) \in \mathbb{C}^{d \times d}$ ,  $\mathsf{B} = (1, \ldots, 1)^{\top} \in \mathbb{C}^{d \times 1}$ ,  $\mathsf{C} = (R_1, \ldots, R_d) \in \mathbb{C}^{1 \times d}$  whose step can be evaluated in  $\mathcal{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} = \mathsf{A}x_{t} + \mathsf{B}u_{t}$  and then take the real part of the output,  $y_{t} = \Re [\mathsf{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]:
1. Given the projections  $q, k, v \in \mathbb{R}^{L \times D}$ , we split them into  $M$  chunks of size  $N = D / M$ ,  $q^m, k^m, v^m \in \mathbb{R}^{L \times N}$ .
2. 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 \coloneqq k^m \otimes v^m \in \mathbb{R}^{L \times N \times N}$ , apply a long convolution with filter  $h^m$  to all  $N \times N$  elements independently, then compute  $y_t^m = (h^m * z^m)_t q_t^m$ ,  $y^m \in \mathbb{R}^{L \times N}$  as shown in Figure 4.
3. Finally, we compose  $y^{1}, \ldots, y^{m}$  into a single output  $y \in \mathbb{R}^{L \times D}$  via concatenation.
Formally, this is related to low-rank approximation characteristics of the Hankel operator; rigorous bounds can be constructed by application of the Eckart-Young-Mirsky theorem [36].
This strategy can also be used to evaluate the filter  $\hat{h}$  alternatively to the standard  $\mathcal{O}(dL)$  method
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].
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.
Figure 4.1: A single head of a multi-head Hyena.
Theorem 4.1. The MultiHyena layer, with  $\mathcal{O}(\log s)$  heads and model size  $\mathcal{O}(\sqrt{s}\log s)$  can solve the associative recall problem, where  $s$  denotes the vocabulary size.
In Appendix E.1, we empirically verify improved scaling in vocabulary size with multiple heads.
- 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: We evaluate 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.
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 300B 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  $\{\lambda_n\}$  of the modal canonical forms (Section 3.2) at different orders  $d$ , and solve for each  $\ell_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
Model	PERP.
GPT	9.3
Hyena	9.3
MultiHyena	8.7
Model	5B	10B	15B
GPT (125M)	13.3	11.9	11.2
Hyena (153M)	13.3	11.8	11.1
MultiHyena (153M)	12.1	11.0	10.6
Model	5B	10B	15B
GPT (355M)	11.4	9.8	9.1
Hyena (355M)	11.3	9.8	9.2
MultiHyena (355M)	10.6	9.4	8.9
Table 5.1: [Left] Perplexity of small models on THE PILE, after pre-training for 300 billion tokens. [Center and Right] Perplexity on THE PILE for models trained until a total number of tokens e.g., 5 billion (different runs for each token total).
Figure 5.1: Errors between logits of pretrained and distilled MultiHyena. In blue, we plot (ordered) logits, in light blue the cumulative distribution function, and in black the relative errors. The green dotted line indicates the  $99.99\%$  percentile. As the errors grow slowly as a function of the percentiles, model outputs do not diverge from the base model.
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 small 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  $\ell_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 ( $>50k$  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, top- $k$ , top- $p$  [43]. Indeed, the relative errors are  $<10^{-2}$  up to and including the  $99.99\%$  percentile of the distribution, meaning e.g., a top- $p$  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.
Model	LAMBADA acc	Winogrande acc	PIQA acc	HellaSwag acc norm.	OpenbookQA acc norm.
Pythia (160M)	32.8	53.1	61.6	31.6	29.2
MultiHyena (154M)	43.2	52.7	64.6	34.1	29.0
LaughingHyena-16	43.1	52.6	64.7	34.1	28.9
LaughingHyena-8	0.0	51.8	51.5	32.7	28.2
LaughingHyena-4	0.0	49.6	53.7	26.4	26.4
Table 5.2: Evaluation of LaughingHyena-distilled models pre and post modal distillation. We test on LM-Eval-Harness [42] and HELM [41] tasks, reporting Pythia [44] performance as a Transformer baseline trained on the same data. LaughingHyena-  $d$  is a MultiHyena model with each filter distilled of order  $d$ .
Figure 5.2: Approximation error profiles (min, max, average) on the filters of MultiHyena model after distillation at different orders. We also visualize the distribution of Hankel singular values: if the spectrum decays after  $n$  singular values, order  $n$  distillation yields low errors.
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\times$  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.
SSM state 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 \times$  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.,  $\tilde{\mathcal{O}}(T)$  as detailed in Section 3.4, compared to  $\mathcal{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).
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.
Figure 5.3: Scaling in prompt  $T$ .
Figure 5.4: Peak GPU memory for generation.
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.
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 deployment, 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].
[51] Luca Perotti and Michal 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).
[52] 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).
[53] Julia Eaton et al. "Polynomial root radius optimization with affine constraints". In: Mathematical Programming 165 (2017), pp. 509-528 (cit. on p. 27).
[54] 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).
[55] 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).
[56] William B Johnson. "Extensions of Lipschitz mappings into a Hilbert space". In: Contemp. Math. 26 (1984), pp. 189-206 (cit. on p. 31).
[57] Devdatt P Dubhashi and Alessandro Panconesi. Concentration of measure for the analysis of randomized algorithms. Cambridge University Press, 2009 (cit. on p. 34).
[58] Ilya Loshchilov and Frank Hutter. "Decoupled weight decay regularization". In: (2017). arXiv: 1711.05101 (cit. on p. 35).
[59] 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).
LAUGHING HYENA DISTILLERY
1 Introduction 1
2 Preliminaries and Related Work 2
2.1 Long Convolution Sequence Models 3
2.2 Auto-Regressive Generation 3
3 The Laughing Hyena Distillery 4
3.1 Data-Free Distillation Objectives 5
3.2 Making Hyena Laugh with Modal Interpolation 5
3.3 Minimal Distillation Orders 6
3.4 Deploying the Recurrence 7
4 Multi-head Long Convolutions 7
5 Experiments 8
5.1 Pre-training 8
5.2 Distillation Analysis 8
5.3 Downstream Evaluation 9
5.4 Benchmarking 10
6 Conclusion 10
A Linear Systems 17
A.1 Extended Notation and System Theory Preliminaries 17
A.2 Systems Norms 17
A.3 Transfer Function of State-Space Models 18
A.4 Truncated Transfer Functions 19
A.5 From Transfer Function to State-Space 20
A.5.1 Isolating the  $h_0$  -term from Transfer Function by Long division 21
A.5.2 Construction of the State-Space from the Transfer Function 21
A.6 From State-Space to Transfer Function 22
A.7 State-Space Representation of Truncated Filters. 23
A.8 Efficient Computation of State-Space Models 23
A.8.1 Fast Evaluation of the Transfer Function 23
A.8.2 Fast Companion Recurrence 24
A.8.3 Canonization of State-Space Models 24
B LaughingHyena: Further Details 26
B.1 Parametrization of Modal Interpolators 26
B.2 Distillation as Rational Interpolation 26
C Proofs 28
C.1 Proof of Lemma 2.1 28
C.2 Proof of Lemma 2.2 28
C.3 Proof of Lemma 2.3 28
C.4 Proof of Proposition 3.1 28
C.5 Proof of Lemma 3.1 29
C.6 Proof of Theorem 3.2 29
C.7 Proof of Proposition 3.3 29
C.8 Proof of Proposition 3.2 30
C.9 Proof of Theorem 4.1 31
D Experimental Details 35
D.1 Pre-training 35
D.2 Distillation Analysis 35
D.2.1 Pretrained Filters: Effective Dimension 35
D.3 Downstream Evaluation 36
D.4 Benchmarking 36
E Additional Experiments 43
E.1 Associative Recall with MultiHyena 43
E.2 Analysis of Hankel Singular Values of Pretrained Large Convolution Sequence Models 43
E.3 Model Order Reduction of H3 43
E.3.1 Modal Truncation 43
E.3.2 Balanced Truncation 44
Authors Contribution