Lizard: An Efficient Linearization Framework for Large Language Models

TL;DR

Lizard is a novel linearization framework that transforms large language models into more efficient, subquadratic architectures with adaptive memory control, maintaining high performance on benchmarks.

Contribution

Lizard introduces a subquadratic attention mechanism with learnable modules and a hardware-aware algorithm to improve efficiency and robustness of large language models.

Findings

Achieves near-lossless performance recovery of teacher models.

Outperforms previous methods by up to 24.5 points on MMLU benchmark.

Demonstrates superior associative recall capabilities.

Abstract

We propose Lizard, a linearization framework that transforms pretrained Transformer-based Large Language Models (LLMs) into subquadratic architectures. Transformers faces severe computational and memory bottlenecks with long sequences due to the quadratic complexity of softmax attention and the growing Key-Value (KV) cache that makes inference memory-bound by context length. Lizard addresses these limitations by introducing a subquadratic attention mechanism that closely approximates softmax attention while preserving model quality. Unlike prior linearization methods constrained by fixed, non-adaptive structures, Lizard augments the architecture with compact, learnable modules that enable adaptive memory control and robust length generalization. Moreover, we introduce a hardwareaware algorithm that solves numerical instability in gated attention to accelerate training. Extensive experiments show that Lizard achieves near-lossless recovery of its teacher model's performance, significantly outperforming previous methods by up to 9.4 - 24.5 points on the 5-shot MMLU benchmark and demonstrating superior associative recall.