PaperlayerView exact revision on arXiv

Source-linked AI summary

Short-Long Convolutions Help Hardware-Efficient Linear Attention to Focus on Long Sequences

Zicheng Liu, Siyuan Li, Li Wang, Zedong Wang, Yunfan Liu, Stan Z. Li

arXiv:2406.08128v3cs.LG

TL;DR

Long-sequence modeling faces quadratic full-attention costs, while existing hybrid approaches have not combined practical linear-attention efficiency with stable convolutional modeling. CHELA pairs hardware-efficient linear attention with short-long convolutions; the paper reports improved benchmark performance and faster, better autoregressive language modeling than the vanilla Transformer.

  • Problem

    Full attention has quadratic time and space complexity, while prior hybrid models are essentially quadratic and their linear variants suffer performance degradation.

    Source: §1 · p. 1

  • Method

    CHELA combines short-long convolutions with hardware-efficient linear attention, using a multilevel convolutional design and an efficient linear-attention implementation.

    Source: §1 · p. 2

  • Results

    CHELA outperforms existing approaches on Long Range Arena and is significantly faster than the vanilla Transformer while yielding better performance in autoregressive language modeling.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The proposed design combines global abstraction with fine-grained data-dependent selection for long-sequence learning.

    Source: §1 · p. 2

    The reported broad evaluation covers image, text, logical reasoning, and speech data.

    Source: §8 · p. 9

Abstract

from arXiv · show

To mitigate the computational complexity in the self-attention mechanism on long sequences, linear attention utilizes computation tricks to achieve linear complexity, while state space models (SSMs) popularize a favorable practice of using non-data-dependent memory pattern, i.e., emphasize the near and neglect the distant, to processing sequences. Recent studies have shown the priorities by combining them as one. However, the efficiency of linear attention remains only at the theoretical level in a causal setting, and SSMs require various designed constraints to operate effectively on specific data. Therefore, in order to unveil the true power of the hybrid design, the following two issues need to be addressed: (1) hardware-efficient implementation for linear attention and (2) stabilization of SSMs. To achieve this, we leverage the thought of tiling and hierarchy to propose CHELA (short-long Convolutions with Hardware-Efficient Linear Attention), which replaces SSMs with short-long convolutions and implements linear attention in a divide-and-conquer manner. This approach enjoys global abstraction and data-dependent selection from stable SSM and linear attention while maintaining real linear complexity. Our comprehensive experiments on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method.

1. Introduction

Full attention’s quadratic complexity makes long-sequence modeling computationally prohibitive, while earlier hybrids remain quadratic or lose performance in linear form. CHELA combines hardware-efficient linear attention with short-long convolutions, and the paper reports strong benchmark and language-modeling results.

1. Introduction

  • Full attention has quadratic time and space complexity, limiting its practicality on long sequences, while prior hybrid models are essentially quadratic and their linear versions degrade in performance.The paper attributes full attention’s overfitting tendency to its lack of structural bias.Source: §1 · p. 1
  • The paper identifies hardware-efficient linear attention and stabilization of long convolution as key challenges for effective hybrid models.It notes that chunk linearization can rival SSM speed but has a significant performance drop, and that long convolution needs regularization to avoid overfitting high-frequency noise.Source: §1 · pp. 1–2
  • CHELA outperforms existing approaches on Long Range Arena and is significantly faster than the vanilla Transformer while yielding better performance in autoregressive language modeling.The paper reports that CHELA outperforms baselines in all evaluated settings.Source: §1 · p. 2
  • CHELA’s linear-attention implementation maintains a constant rate regardless of sequence length with fixed memory.The contribution is presented as releasing linear attention’s speed potential in attention-SSM hybrids.Source: §1 · p. 2
  • CHELA combines linear attention with short-long convolutions to provide global abstraction and fine-grained, data-dependent selection.The paper describes the convolutional hierarchy as improving multi-frequency learning, performance, and stability for long convolution.Source: §1 · p. 2

2. Background

The background introduces attention’s pairwise token relations and a mathematically equivalent linear variant. It also describes recurrent and convolutional formulations, including the difficulty of efficiently determining a convolution kernel.

2.1. Self Attention

  • For self-attention, the matrix M records pairwise token relations, and linear attention transforms the score computation through right matrix multiplication while remaining mathematically equivalent.M_ij represents the relation between input tokens i and j.Source: §2.1 · p. 2

2.2. Linear Recurrent Models

  • A continuous-time linear recurrent model maps a one-dimensional input u(t) to a latent state x(t), then maps that state to a one-dimensional output y(t).The state has d_s dimensions; the model can be discretized for discrete sequences such as natural-language inputs.Source: §2.2 · p. 3
  • HiPPO matrices initialize A so the state x(t) can remember the input history up to the present.The passage attributes this history-retention property to the construction of the matrices.Source: §2.2 · p. 3
  • The recurrent representation can be written as a convolution y = K * u, and the kernel can be computed efficiently using O(L) time and space.The passages describe the convolution kernel as a sequence of terms involving C, A, and B.Source: §2.2 · p. 3
  • When the convolution kernel K is known, output computation can be efficient, but determining K is difficult and most existing algorithms require O(L^2) time and space.A learned kernel whose length matches the sequence can replace an SSM layer as a drop-in replacement.Source: §2.2 · p. 3
  • The S4 model is presented as a way to calculate the recurrent model’s convolution formulation, with A and B initialized using a specified construction.The supplied passages identify Gu et al. (2021) as devising S4 to calculate the equation and describe the initialization context.Source: §2.2 · p. 3

L DKFLu (7)

The passage identifies the DFT matrix and a diagonal matrix derived from the kernel as components of FFT convolution, which scales linearly with sequence length.

L DKFLu (7)

  • FFT convolution uses the DFT matrix F_L and D_K = diag(F_LK), and scales linearly with sequence length L.The passage defines F_L as the DFT matrix of size L.Source: L DKFLu (7) · p. 3

3. Why Combining Long Conv and Attention?

Hybrid long-sequence modeling combines structured memory with data-dependent patterns, but practical use is constrained by linear attention’s hardware bottlenecks and long convolution’s stability needs.

3. Why Combining Long Conv and Attention?

  • CHELA combines structured-memory inductive bias with data-dependent patterns to model long sequences in true linear time.Figure 2 depicts the CHELA layer and its Short-Long Convolutions and Flash Linear Attention components.Source: §3 · p. 3, Figure 2 · p. 4
  • Linear attention’s practical speed is limited by intensive GPU memory access and cumulative summation in causal settings, despite its O(L) complexity.CHELA’s hardware-aware design targets these bottlenecks.Source: §3 · p. 3
  • A single long-convolution kernel struggles to learn multiple frequencies, motivating multiresolution short kernels that share its learning pressure and can be fused for inference.Long convolution also needs additional regularization for stable performance across data types.Source: §3 · p. 3

4. CHELA: Short-Long Convolutions with Hardware-Efficient Linear Attention

CHELA combines hardware-aware linear attention with short-long convolutions to address practical efficiency and stability challenges in long-sequence modeling. Its design uses tiled computation, multiresolution kernels, gating, and structural reparameterization.

4.1. Hardware-Efficient Linear Attention

  • Tiling computes linear attention in SRAM using divide-and-conquer loops, accumulating intermediate results there before returning outputs to HBM.The approach targets the theoretical speed of linear attention by exploiting GPU memory-bandwidth differences.Source: §4.1 · p. 4

4.2. Short-Long Convolutions

  • Short convolutions before a long kernel target multi-frequency learning and improve long-convolution stability and capacity.The short kernels are fused into a single kernel by structural reparameterization for inference.Source: §4.2 · pp. 4–5
  • On ListOps, the short-kernel design demonstrates better stabilization and accuracy for long-sequence logic inference.The experiments use a fixed kernel size of three and a length-dependent size of 2 log10 L + 1.Source: Figure 3 · p. 4, §4.2 · p. 5

4.3. CHELA Model

  • CHELA embeds short-long convolution output Z into gated linear attention, then uses an output gate to update the final activation in a residual style.The block also includes a two-layer feedforward network, layer normalization, and residual connections, with pre-layer normalization.Source: §4.3 · p. 5

4.4. Relation to MEGA and SPADE.

  • CHELA is presented as more efficient and higher-performing than MEGA and SPADE, with linear-time speedup and structured representations enhanced by short convolutions.The comparison notes chunking-related performance limits in the linear variants of MEGA and SPADE.Source: §4.4 · p. 5

5. Experiments

Experiments evaluate CHELA on long-context, raw-speech, language-modeling, and pixel-sequence tasks. The reported results include strong benchmark performance and faster inference or processing in several settings.

5. Experiments

  • Experiments used NVIDIA A100-80G and PyTorch, with float32 parameters and bfloat16 for most computations.The evaluation covered five standard sequence-modeling tasks and compared against recent models for each task.Source: §5 · p. 5

5.1. Long-Context Sequence Modeling

  • 88.26% average accuracy across six LRA tasks surpassed all compared methods; byte-level classification with a 4K input was 5.8 times faster.The LRA inputs range from 1K to 16K tokens; results average five runs.Source: §5.1 · p. 6

5.2. Raw Speech Classification

  • 97.98% accuracy with 493K parameters made CHELA the leading method in the Speech Commands comparison.The experiment classifies unaltered speech signals of duration 16000.Source: §5.2 · p. 6

5.3. Auto-Regressive Language Modeling

  • CHELA outperformed baseline models on WikiText-103 and enwik8, with inference almost 10 times faster than the Pure Transformer.The authors report that the short-long convolution layer manages length extrapolation during inference.Source: §5.3 · p. 7

5.4. Pixel-Level Sequential Image Classification

  • CHELA achieved the best test accuracy on sequential CIFAR-10, exceeding previous convolution-based models by more than ten percentage points.The most effective model used ten CHELA blocks.Source: §5.4 · p. 7

6. Ablation Study

The ablations test hardware-efficient linear attention and short-long convolutions. The reported comparisons find near-real-linear scaling and identify short-long convolutions as the strongest linear-attention partner.

6. Ablation Study

  • The ablation study centers on whether hardware-friendly implementation improves linear-attention speed and whether short-long convolutions help on long sequences.Table 6 is identified as an ablation study of structured mixers in CHELA.Source: §6 · p. 7, Table 6 · p. 8

Q1: Benchmark hardware-efficient linear attention

  • Hardware-efficient linear attention showed almost real-linear scaling with sequence length and more than doubled the speed of the original PyTorch implementation.The comparison used WikiText-103 models with 200M parameters; runtime is visualized in Figure 4.Source: §6 · p. 8

Q2: Analysis of short-long convolutions

  • Among the tested SSM-like modules on LRA Text, Image, and PathX, short-long convolutions were reported as the best partner for linear attention.The comparison combined different hybrid models following CHELA's modeling structure.Source: §6 · p. 8

7. Related Works

Related work covers approaches to efficient attention, state-space models and long convolutions, and hardware-efficient implementation. It describes trade-offs in approximation, initialization, and memory movement.

Efficient transformer models

  • Sparse attention limits each token to predefined subsets, while clustering methods partition queries or keys for inter-cluster attention.Kernel-based methods approximate full attention in linear time, but the cited discussion notes significant approximation errors for kernel and low-dimensional projection methods.Source: §7 · p. 8

State space models and long convolutions

  • S4 can be implemented through convolutions, but its input-length long kernel requires sophisticated parameterization and approximation, and initialization can require special effort.The passage also notes models based on linear recurrence or global convolution.Source: §7 · p. 8

Hardware-efficient implementation

  • FlashAttention optimizes standard attention on GPUs using tiling to reduce memory reads and writes between HBM and on-chip SRAM.Source: §7 · p. 8

8. Conclusion and Limitations

CHELA combines short-long convolutions with hardware-efficient linear attention to model long sequences in linear time. The authors report broad performance and efficiency, while noting that short-convolution combinations remain unexplored.

8. Conclusion and Limitations

  • CHELA computes structured and data-dependent patterns in linear time using an SSM-attention hybrid architecture.The authors attribute its performance to the global view of short-long convolutions and hardware-efficient linear attention with gating.Source: §8 · pp. 8–9
  • CHELA reports performance improvement over other hybrid linear models while maintaining linear complexity, with experiments spanning image, text, logical reasoning, and speech.Source: §8 · p. 9
  • The optimal combinations of short convolutions were not explored; designing dynamic short-convolutional components according to input data is future work.The authors also identify time-varying SSMs as a different idea for this goal.Source: §8 · p. 9

Impact Statement

The authors describe the work's intended impacts and ethical scope, and encourage early discussion of potential negative impacts as new methods emerge.

Impact Statement

  • The paper aims to advance long-sequence modeling with an efficient design and make sequence mixers more efficient by leveraging hardware and different model features.Source: Impact Statement · p. 9
  • The authors state they do not foresee direct negative societal consequences and used public resources without private or sensitive information.Source: Impact Statement · p. 9
  • The authors encourage discussing potential negative impacts early as the community proposes new methods.Source: Impact Statement · p. 9

A. Experimental Details

The appendix describes data and configuration choices for Long Range Arena and sCIFAR, then language modeling on WikiText-103 and enwik8.

A.1. Long Range Arena (LRA) and sCIFAR

  • LRA preprocessing and data splits follow Tay et al. (2020b), and sCIFAR uses the image-task parameter settings from LRA.Source: §A.1 · p. 13
  • Table 7 lists CHELA hyperparameters for LRA and raw speech classification, with BSZ, LR, and WD denoting batch size, learning rate, and weight decay.BN, LN, and SN denote Batch, Layer, and Scale Normalization.Source: Table 7 · p. 13

A.2. Language Modeling

  • Language modeling uses WikiText-103 and enwik8 data and splits from Ma et al. (2022), with training segments formed from m consecutive chunks.Chunk size is the effective attention length; other hyperparameters are presented in Table 8.Source: §A.2 · p. 13, Table 8 · p. 13
Loading 2406.08128v3…