PaperlayerView exact revision on arXiv

Source-linked AI summary

Linear Transformers with Learnable Kernel Functions are Better In-Context Models

Yaroslav Aksenov, Nikita Balagansky, Sofia Maria Lo Cicero Vaina, Boris Shaposhnikov, Alexey Gorbatovski, Daniil Gavrilov

arXiv:2402.10644v2cs.LGcs.CL

TL;DR

Long-context Transformer attention is computationally costly, and existing linear architectures still face kernel and retrieval challenges. ReBased makes the Based kernel learnable and adds normalization, with reported gains over Based on MQAR and Pile language modeling.

  • Problem

    Transformer attention scales poorly on long sequences, and the Based kernel is considered suboptimal for lengthy contexts and small model capacities.

    Source: §1 · p. 1

  • Method

    ReBased uses a polynomial kernel with learnable parameters and integrates normalization directly into the kernel function.

    Source: §6 · p. 8

  • Results

    ReBased outperforms Based on MQAR and Pile language-modeling tasks, including in-context learning and associative-dependency perplexity.

    Source: §1 · p. 2, §6 · p. 8

  • Takeaways & Limitations

    The findings suggest that more sophisticated learnable kernels and normalization may improve kernel-based models.

    Source: §6 · p. 8

    The experiments use academic-scale models, limiting how confidently findings can be extrapolated to larger models.

    Source: §7 · p. 9

Abstract

from arXiv · show

Advancing the frontier of subquadratic architectures for Language Models (LMs) is crucial in the rapidly evolving field of natural language processing. Current innovations, including State Space Models, were initially celebrated for surpassing Transformer performance on language modeling tasks. However, these models have revealed deficiencies in essential In-Context Learning capabilities - a domain where the Transformer traditionally shines. The Based model emerged as a hybrid solution, blending a Linear Transformer with a kernel inspired by the Taylor expansion of exponential functions, augmented by convolutional networks. Mirroring the Transformer's in-context adeptness, it became a strong contender in the field. In our work, we present a singular, elegant alteration to the Based kernel that amplifies its In-Context Learning abilities evaluated with the Multi-Query Associative Recall task and overall language modeling process, as demonstrated on the Pile dataset.

1 Introduction

Transformers handle in-context learning well but scale poorly to long sequences, while linear alternatives and the Based model leave performance and kernel-design challenges. The paper introduces ReBased, which modifies the kernel to improve long-context retrieval and reports stronger MQAR and Pile results than Based.

  • Transformer attention scales poorly with long sequences, making extended contextual processing computationally impractical.Source: §1 · p. 1
  • Linear Transformers and State Space Models can be described as RNNs, whose memory capacity can be overrun as information accumulates; SSMs nevertheless showed higher text-modeling quality than Linear Transformers.Source: §1 · p. 1
  • The Based kernel is argued to be suboptimal for lengthy contexts and small model capacities.Source: §1 · p. 1
  • ReBased modifies the kernel so attention can assign specific tokens zero probability, improving long-sequence retrieval while simplifying attention computation.Source: §1 · p. 2
  • ReBased surpasses Based across varied MQAR contexts and model sizes, and performs better on Pile in in-context learning and associative-dependency perplexity.Figure 1 reports that ReBased outperforms all baselines except Attention across sequence lengths and model sizes.Source: §1 · p. 2, Figure 1 · p. 2

2 Recent Work

Recent work seeks alternatives to quadratic-cost Transformer attention, but long-context in-context learning remains challenging, especially for small models. This paper targets that remaining performance decline.

  • Vanilla Transformer computation and memory grow as O(d*N^2), becoming a bottleneck for long sequences.Source: §2 · p. 2
  • Kernel-based linear attention changes computation order and reduces time and memory complexity to O(d^2*N), with inference sampling linear in sequence length.Source: §2 · p. 2
  • Perplexity may not reflect long-context dependencies, while Long Range Arena comparisons may be biased; MQAR instead tests retrieval of multiple tokens from context.Source: §2 · pp. 2–3
  • Based outperforms Mamba on MQAR under long-context and small-model conditions, yet still declines relative to Transformers on long sequences with smaller models.Source: §2 · p. 3
  • The paper's primary goal is to address the decline of smaller models on long sequences relative to traditional Transformers.Source: §2 · p. 3

3 Background

Linear Transformers replace direct query-key similarity with kernel features to reduce sequence-length complexity. Based uses a Taylor-inspired kernel and convolutional layers, improving on Mamba in some MQAR settings but not eliminating the gap to Transformers.

3.1 Linear Transformers

  • The original Transformer's exponential query-key similarity entails quadratic time and memory growth as sequences lengthen.Source: §3.1 · p. 3
  • Linear Transformer kernelization approximates query-key similarity with feature products and rearranges computation to make attention linear in sequence length.Source: §3.1 · p. 3

3.2 Based

  • Based uses a Taylor-inspired kernel that approximates the exponential over a specific query-key range, alongside convolutional layers for short non-associative dependencies.Source: §3.2 · p. 3
  • Based outperforms Mamba on MQAR for substantial contexts and constrained model capacity, but still trails the original Transformer.Source: §3.2 · p. 3

4 Revisiting Based

Based's quadratic kernel has a fixed positive minimum, making it difficult to suppress attention to selected token pairs. ReBased makes the kernel's shape trainable and adds normalization before kernel evaluation.

  • The Based kernel's minimum similarity is fixed at 0.5, making near-zero attention scores difficult for selected token pairs in long sequences.Source: §4 · p. 3, Figure 2 · p. 4
  • ReBased makes the quadratic kernel adjustable so its parabola can align with the range of query-key products encountered during training.Source: §4 · p. 4
  • Separate query and key parameters let the model learn non-negative quadratic kernels with a single real root.Source: §4 · p. 4
  • The model also applies Layer Normalization before the kernel function, which the authors report improves results in an ablation study.Source: §4 · p. 4

5 Experiments

Across MQAR, Pile language modeling, and few-shot benchmarks, ReBased generally improves on Based. The attention model retains advantages in some settings, including MQAR attention overlap and in-context token attending.

5.1 Experimental Setup

  • MQAR evaluations covered sequence lengths 128–2048 and hidden sizes 64–512, while Pile language modeling used sequence length 4096 and separated AR from Non-AR perplexity.The MQAR architecture used a short convolution in the first layer and a specified mixer in the second.Source: §5.1 · pp. 4–5

5.2 MQAR experiment

  • At sequence length 2048, ReBased matched or surpassed Based as model dimensions increased, while RWKV and Mamba failed across tested sizes.At dimension 64, all models except Attention struggled; the performance gap between Attention and the other models becomes pronounced above sequence length 512.Source: §5.2 · p. 6

5.3 Ablation Study

  • ReBased performs best on all model dimensions; differences are not very significant at sequence length 256 but become more evident at 2048, while normalization before the kernel improves over Based.Affine transformations show favorable performance compared to the x2 model, which does not significantly decrease with sequence length.Source: Figure 3 · p. 5, §5.3 · pp. 6–7

5.4 Language Modeling

  • On Pile perplexity, ReBased outperformed Based on AR and Non-AR tokens, but Attention retained an AR-perplexity advantage and the strongest in-context attending ability.ReBased outperformed both Based architectures on Non-AR perplexity, yielding better overall perplexity than Based.Source: Table 2 · p. 7, §5.4 · p. 7

5.5 Few-Shot Performance

5.6 Analysis

  • Attention achieved higher IoU than Based and ReBased, while ReBased narrowed the performance divide between sub-quadratic methods and attention.The analysis compares attention matrices with the ground-truth retrieval positions on MQAR.Source: §5.6 · p. 7, Table 5 · p. 8, Figure 4 · p. 8

6 Conclusion and Future Work

ReBased uses polynomial kernels with learnable parameters and normalization before kernel evaluation, and outperforms Based on MQAR and Pile language-modeling tasks. Its representations resemble attention more than Based, while the findings point to unresolved long-sequence gaps and opportunities for more adaptable kernels.

6 Conclusion and Future Work

  • ReBased's learnable polynomial kernel and normalization before kernel evaluation outperform Based on MQAR and Pile language-modeling tasks.The paper describes normalization as integrated directly into the kernel function.Source: §6 · p. 8
  • ReBased representations resemble attention more than Based, and its kernel differs from the exponential-inspired kernel used by Based.The paper suggests second-order polynomials may be insufficient and that more sophisticated learnable kernels could improve trained models.Source: §6 · p. 8
  • Attention models significantly outperform nonattention counterparts on MQAR as sequence lengths increase, leaving a gap for future research to bridge.The paper notes associative recall as a relevant task area, including machine translation.Source: §6 · p. 8

7 Limitations

ReBased applies to a wide range of tasks typically addressed by Transformers, but its effectiveness on intensive copying or recall remains unclear. Experiments are also limited to academic-scale models, constraining extrapolation to larger models.

7 Limitations

  • ReBased is presented as applicable to a wide range of tasks typically addressed by Transformers.The paper says it remains unclear how effective the method is on tasks involving intensive copying or recalling previous context.Source: §7 · p. 9
  • Effectiveness remains unclear for tasks involving intensive copying or recalling previous context.The paper says addressing these tasks is crucial for fully mitigating inference problems associated with attention mechanisms.Source: §7 · p. 9
  • Experiments are limited to academic-scale models, so the findings' extension to larger models is uncertain.The authors identify resource constraints as the reason for the scale boundary.Source: §7 · p. 9

A Details for the MQAR dataset experiments

The MQAR experiments use adjusted, more challenging data configurations and report results across five seeds. The appendix also describes architecture modifications and shared training settings.

A Details for the MQAR dataset experiments

  • Replacing the first attention layer is beneficial for performance, while RWKV performs better without that replacement.The main experiment therefore uses two RWKV layers; the passage also notes a modified attention model chosen for stability.Source: A Details for the MQAR dataset… · p. 11, Figure 6 · p. 11
  • The data configurations were adjusted to make the task more challenging, with sequence lengths and QK-pair counts documented in Table 7.The table caption identifies its contents but does not provide the values in these passages.Source: A Details for the MQAR dataset… · p. 11, Table 7 · p. 11
  • All experiments use five seeds and batch size 512, with gradient accumulation when GPU memory is insufficient.Learning rates were selected from 5e-4, 1e-3, 3e-3, and 1e-2.Source: A Details for the MQAR dataset… · p. 11

B Stability

ReBased reaches 100% accuracy across the tested hyperparameter sets and seeds when model dimension can store the necessary information, while Based degrades at some learning rates. The appendix assesses stability with an Expected Validation Performance plot and averages accuracy across five seeds.

B Stability

  • The stability analysis uses an Expected Validation Performance plot and treats the five-seed average as final accuracy.The plot is presented in Appendix Figure 7.Source: B Stability · p. 11
  • 100% accuracy for every tested hyperparameter set and seed was observed for ReBased when model dimension could store all necessary information; Based degraded at some learning rates.The result is reported as evidence of greater training stability for ReBased.Source: B Stability · pp. 11–12
  • The reported 100% result applies even when model dimension is sufficiently large to store all necessary information.The passage contrasts this condition with Based's degradation at certain learning rates.Source: B Stability · p. 12

C Pile Dataset Experiment Details

The Pile experiment uses a tokenization different from Based’s and omits rotary positional embeddings and other tricks. Based/ReBased use short convolutions as mixers in every odd-numbered layer.

C Pile Dataset Experiment Details

D Additional Analysis

Additional analyses compare attention patterns across models and examine hyperparameters and layer-normalization parameters. ReBased is reported to perform best across hyperparameters, budgets, and model dimensions, while its attention remains noisier than the Attention baseline.

D Additional Analysis

  • ReBased has fewer noisy tokens than Based, but Attention has superior attention scores and a peak at the ground-truth position.The attention-matrix analysis measures IoU against ground-truth positions over 10,000 examples; results are presented in Table 5.Source: D Additional Analysis · p. 12, Figure 8 · p. 13, Figure 9 · p. 14
  • Layer-normalization parameters converge independently of training sequence length, while gamma and beta have high standard deviations relative to their mean absolute values.The authors suggest this may indicate a need to provide features with different scales; Figure 10 shows gamma tending to about 0.13 and beta to 0.Source: D Additional Analysis · p. 12, Figure 10 · p. 14
  • ReBased works best across all tested hyperparameters, budgets, and model dimensions.Figure 7 presents expected validation accuracy across these settings.Source: Figure 7 · p. 13

Related papers

Loading 2402.10644v2…