[Paper Review] Laughing Hyena Distillery: Extracting Compact Recurrences From Convolutions

Stefano Massaroli, Michael Poli+12 more|VU Research Portal|Oct 28, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper introduces Laughing Hyena, a distillation method that extracts compact, recurrent state-space models (SSMs) from pre-trained long convolution sequence models (LCSMs), enabling O(1) compute and memory per token during autoregressive generation. The approach achieves 10× higher throughput than Transformers and 1.5× higher than Hyena at 1.3B parameters, with no loss in quality after distillation.

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.

Motivation & Objective

• To enable constant-time, constant-memory autoregressive generation in long convolution sequence models (LCSMs), overcoming the O(K) memory and O(K²) compute bottleneck of standard inference.
• To develop a distillation framework that extracts low-dimensional, stable state-space models (SSMs) from pre-trained convolutional layers while preserving model quality.
• To improve pre-training quality and distillation efficiency by re-designing the Hyena architecture with weight-tied filters across channels.
• To enable high-throughput, large-batch generation by reducing memory footprint and enabling recurrent inference in LCSMs.

Proposed method

• Applies rational interpolation and model-order reduction to extract compact SSMs from convolution filters of pre-trained LCSMs.
• Introduces a factorized modal parametrization inspired by barycentric and Prony-like methods to improve SSM stability and avoid numerical issues.
• Uses discrepancy metrics on convolution filters as the approximation objective, ensuring compatibility across downstream tasks.
• Employs Hankel operator spectral analysis to determine optimal state dimension d for SSMs.
• Revises the Hyena block with weight tying across channels to increase effective filter dimension and reduce distillation complexity.
• Employs balanced truncation and modal truncation for SSM model reduction, with error bounds derived from Gramian eigenvalues.

Experimental results

Research questions

• RQ1 Can compact, recurrent SSMs be extracted from pre-trained long convolution models to enable O(1) inference cost per token?
• RQ2 What state dimension d yields optimal trade-off between approximation error and model efficiency in SSM distillation?
• RQ3 How can SSM parametrization be improved to avoid numerical instability and improve convergence during distillation?
• RQ4 Can architectural modifications to the Hyena block enhance pre-training quality and reduce distillation cost?
• RQ5 Does distillation of LCSMs into SSMs preserve downstream performance while enabling high-throughput generation?

Key findings

• Laughing Hyena achieves 10× higher peak throughput than comparable Transformers and 1.5× higher than Hyena at 1.3B parameters, with no loss in quality after distillation.
• At 1.3B parameters, Laughing Hyena generates 512 tokens using 3× less memory than a Transformer under the same memory constraint.
• The model enables constant O(d) memory and O(dK) time for K-token generation, compared to O(K) and O(K²) for kv-cached Transformers.
• Weight tying across channels in the Hyena block improves pre-training perplexity and reduces the number of filters to be distilled.
• Balanced truncation and modal truncation methods show non-monotonic error reduction in some layers, indicating sensitivity to numerical stability.
• The use of filter discrepancy metrics as the approximation objective ensures robustness across diverse downstream tasks.