PaperlayerView exact revision on arXiv

Source-linked AI summary

A Study on ReLU and Softmax in Transformer

Kai Shen, Junliang Guo, Xu Tan, Siliang Tang, Rui Wang, Jiang Bian

arXiv:2302.06461v1cs.CLcs.AIcs.LG

TL;DR

The paper examines how differing ReLU and Softmax activations affect the relationship between Transformer components and key-value memory. It studies these activations and develops a fully-ReLU Transformer, reporting better performance on long-document translation and identifying scope and latency limits.

  • Problem

    FFNs and key-value memories have similar formulations but use different activations, and their relationship under those choices needs examination.

    Source: §1 · p. 1, §2.2 · p. 3

  • Method

    The paper studies ReLU and Softmax in FFNs, memory, and self-attention, then uses variance-control designs to build a fully-ReLU Transformer.

    Source: §1 · p. 1, §4.2 · p. 6, §6 · p. 9

  • Results

    ReLU performs better with many memory slots and on long inputs, and the full ReLU architecture significantly outperforms vanilla Transformer in long-document translation.

    Source: §3.5 · p. 5, §4.4 · p. 9, §6 · p. 9

  • Takeaways & Limitations

    FFNs and self-attention can be viewed as key-value memories, with layer normalization making FFNs and key-value memory equivalent.

    Source: §3.5 · p. 5

    Experiments are limited to machine translation, and ReLU yields only slight latency gains because self-attention remains O(N^2).

    Source: §7 · p. 9

Abstract

from arXiv · show

The Transformer architecture consists of self-attention and feed-forward networks (FFNs) which can be viewed as key-value memories according to previous works. However, FFN and traditional memory utilize different activation functions (i.e., ReLU and Softmax respectively), which makes them not equivalent. In this paper, we first rebuild the connections between FFN and key-value memory by conducting extensive studies on ReLU and Softmax, and find they are equivalent when adding an additional layer normalization module on Softmax. In addition, ReLU outperforms Softmax on both FFN and key-value memory when the number of value slots is large. We analyze the reasons and then explore this good property of ReLU on the self-attention network where the original Softmax activation performs poorly on long input sequences. We then propose a full ReLU architecture named ReLUFormer which performs better than the baseline Transformer on long sequence tasks such as document translation. This paper sheds light on the following points: 1) Softmax and ReLU use different normalization methods over elements which lead to different variances of results, and ReLU is good at dealing with a large number of key-value slots; 2) FFN and key-value memory are equivalent, and thus the Transformer can be viewed as a memory network where FFNs and self-attention networks are both key-value memories.

1. Introduction

The paper revisits Transformer components as key-value memories while accounting for their different activation functions. It connects FFNs and memory, analyzes ReLU in self-attention, and proposes a fully-ReLU Transformer for long-document translation.

1. Introduction

  • Earlier work connected FFNs and key-value memories by their formats, but the paper highlights their differing activations: ReLU in FFNs and Softmax in memory.The paper's studies revisit these connections while accounting for activation functions.Source: §1 · p. 1
  • FFNs and key-value memory are equivalent when Softmax is followed by layer normalization, while ReLU performs better with larger numbers of value slots.The paper attributes Softmax's performance drop to normalized activations becoming too small to carry enough information.Source: §1 · pp. 1–2
  • Softmax can concentrate weight on a few slots and neglect others, whereas ReLU can use more slots; ReLU's self-attention output variance can also grow with sequence length.The authors introduce a variance reduction factor and regularization losses to address unstable training and report better performance than Softmax on long inputs.Source: §1 · p. 2
  • ReLUFormer uses ReLU in both self-attention and FFN and outperforms the Transformer baseline on long-document translation.The architecture is presented as global key-value memories in FFNs and local memories in self-attention.Source: §1 · p. 2

2. Background

The background describes FFNs and self-attention in relation to key-value memory. It identifies activation functions as a difference that motivates the paper's comparisons.

2.1. Feed-Forward Network and Key-Value Memory

  • An FFN applies two linear projections with a nonlinear activation; earlier work treats its weights as keys and values in a memory.The paper revisits this relation because FFNs typically use ReLU while key-value memory uses Softmax.Source: §2.1 · p. 2

2.2. Self-Attention Network and Key-Value Memory

  • Key-value memory aggregates value parameters using a distribution computed from keys, and self-attention has the same key-value computation form.Self-attention can be viewed as local memory, with the current token as query and other context tokens as keys and values.Source: §2.1 · p. 2, §2.2 · p. 3
  • Although FFN, self-attention, and key-value memory have similar formulations, their activation-function differences had not been carefully discussed.The background frames the paper's analyses of ReLU and Softmax as addressing that gap.Source: §2.2 · p. 3

3. Connections Between FFN and Key-Value Memory

ReLU and Softmax produce different FFN behavior, but layer-normalized Softmax restores FFN–memory comparability. ReLU’s advantage grows with slot count, which the authors relate to Softmax’s increasingly concentrated distribution.

3.1. ReLU and Softmax are Different

  • 34.22 to 33.08 BLEU: replacing FFN with Softmax key-value memory reduces performance on IWSLT14 De-En.The experiment replaces FFN layers with key-value memories while keeping other Transformer components the same.Source: §3.1 · p. 3

3.2. Bridge the Gap between FFN and Key-value Memory

  • Layer normalization after Softmax raises its variance ratio and BLEU, making FFN with Softmax perform similarly to ReLU and restoring FFN–memory equivalence.Without layer normalization, Softmax’s smaller output variance lets the residual dominate the FFN output, reducing parameter utilization.Source: §3.2 · p. 3

3.3. Scaling to Large Number of Values

  • At memory sizes 3072 and 4096, ReLU outperforms Softmax with layer normalization; across tested sizes, Softmax with layer normalization is comparable to ReLU.The study varies the FFN hidden dimension, viewed as the number of memory slots, from 32 to 4096.Source: Figure 1 · p. 3, §3.3 · p. 4

3.4. Quantitative Analysis between ReLU and Softmax

  • Softmax’s top 0.2% of scores occupy more than 85% of the total, with concentration worsening as memory grows; most slots consequently receive too little gradient.The authors conjecture that exponential normalization centralizes Softmax scores, whereas ReLU can aggregate information from more slots.Source: §3.4 · p. 4
  • Softmax value slots collapse in anisotropy analysis, layer normalization alleviates the collapse, and ReLU yields the best result.Anisotropy is the average pairwise similarity between value slots; lower scores indicate more distinct slots.Source: Figure 3 · p. 4, §3.4 · p. 5

3.5. Summary

  • Softmax’s centralized distribution can underuse large memories, while layer normalization can make it comparable to FFN; the authors connect broader slot use to long-sequence self-attention.The summary attributes ReLU’s larger variance to greater expressiveness and identifies Softmax’s centralized distribution as a concern when memory is large.Source: §3.5 · p. 5

4. ReLU vs Softmax in Self-Attention

This section investigates ReLU in self-attention, where sequence length determines the number of memory slots, and develops techniques to stabilize and regularize ReLU attention. ReLUFormer performs better than Transformer and Sparsemax baselines on long-document translation, while remaining comparable at shorter lengths.

4.1. Solving Variance Exploding Caused by ReLU

  • ReLU SAN output variance grows with sequence length, making direct Softmax replacement unstable; adding layer normalization alone still yields sub-optimal BLEU.The paper attributes the variance growth to the sequence-length-dependent output distribution.Source: §4.1 · p. 5
  • A sequence-length-dependent variance reduction factor stabilizes ReLU SAN; the resulting model scores 33.19 BLEU on IWSLT14 De-En, 1.03 BLEU below vanilla Softmax attention.The factor is adapted during end-to-end training because sequence length is dynamic.Source: §4.1 · p. 6

4.2. Closing the Gap Between ReLU and Softmax

  • ReLU attention has very low entropy, with 94% zero weights; normalization and entropy-margin regularization encourage more nonzero weights and constrain entropy.The paper reports that adding the loss makes the distribution flatter and more informative.Source: §4.2 · p. 6

4.3.1. EXPERIMENTS ON SENTENCE-LEVEL TRANSLATION

  • On sentence-level translation, ReLUFormer outperforms Transformer, Sparsemax, 1.5Entmax, and ReLA, with inference latency comparable to Transformer and 1.7 times faster than Sparsemax and 1.5Entmax.The reported comparisons cover IWSLT14 De-En and WMT14 En-De.Source: §4.3.1 · p. 7, Table 4 · p. 7
  • Removing the scale factor prevents convergence, while removing regularization reduces BLEU by 1.37 and entropy by 1.52.The authors interpret the regularization result as evidence that it provides a more informative attention distribution.Source: §4.3.1 · p. 7, Table 5 · p. 7

4.3.2. EXPERIMENTS ON DOCUMENT-LEVEL TRANSLATION

  • At lengths 512, 1024, and 2048, ReLUFormer consistently outperforms Transformer and Sparsemax; at length 1024, it gains 1.15 BLEU, while Sparsemax fails to converge on long sequences.At lengths 128 and 256, ReLUFormer is comparable with both baselines.Source: Table 6 · p. 8, §4.3.2 · p. 8
  • For 1024-token documents, Softmax attention is more centralized than ReLU, while the visualized ReLU map captures more distant correlations and assigns less attention to some stop words.The authors connect Softmax's centralized distribution with attending to fewer contexts on long sequences.Source: §4.3.2 · p. 8, Figure 4 · p. 8, Figure 5 · p. 8

4.4. Summary

  • The section concludes that ReLU's sequence-length-dependent variance requires stabilization, while its less centralized attention distribution supports better long-sequence translation than Softmax.The authors also report that their techniques make ReLU competitive with Softmax on sentence-level translation.Source: §4.4 · pp. 8–9

5. Insights and Findings

The paper distinguishes ReLU and Softmax by their variance and normalization behavior, and finds ReLU useful with many value slots and long sequences. With layer normalization, FFNs and key-value memory are equivalent; Transformer components can be viewed as global and local memories.

5. Insights and Findings

  • ReLU produces larger variance than Softmax, while ReLU self-attention variance grows with sequence length and can cause variance explosion.Small-variance outputs can be dominated by the residual, while sequence-dependent variance can destabilize self-attention.Source: §5 · p. 9
  • Softmax concentrates scores more than ReLU, restricting context use across many value slots; ReLU is less centralized and handles large slot counts and long sequences.The paper connects ReLU’s less centralized distribution with integrating information from more tokens in long sequences.Source: §5 · p. 9
  • Adding layer normalization to Softmax makes key-value memory equivalent to the FFN, while alleviating Softmax’s scale and over-centralization problems.Source: §5 · p. 9
  • The Transformer can be viewed as a memory network: FFNs are global memories with shared parameter keys and values, while self-attention uses keys and values constructed locally from input sequences.Source: §5 · p. 9

6. Conclusion

The paper concludes that layer normalization makes FFNs and key-value memory equivalent, while ReLU performs better with many memory value slots. Its full-ReLU architecture is competitive in sentence-level translation and outperforms the vanilla Transformer in long-document translation.

6. Conclusion

  • The FFN and key-value memory are equivalent when layer normalization is introduced.Source: §6 · p. 9
  • ReLU performs better than Softmax when the number of memory value slots is large.Source: §6 · p. 9
  • With specific designs, the full-ReLU architecture is competitive in sentence-level translation and significantly outperforms the vanilla Transformer in long-document translation.Source: §6 · p. 9

7. Limitation and Future Works

The paper’s experiments are limited to machine translation, and replacing Softmax with ReLU yields only slight latency gains because self-attention remains O(N^2). Future work includes additional tasks and methods to improve self-attention latency.

7. Limitation and Future Works

  • Experiments in this work cover only machine translation.Source: §7 · p. 9
  • Replacing Softmax with ReLU yields slight latency gains because self-attention remains O(N^2).Source: §7 · p. 9
  • The authors plan experiments on tasks such as language modeling and summarization, and further work on self-attention methods to improve latency.Source: §7 · p. 9

A.1. Datasets

The sentence-level machine translation evaluation uses IWSLT14 De-En and WMT14 En-De. The appendix specifies their data splits and shared subword-processing setup.

A.1. Datasets

  • The evaluation uses IWSLT14 De-En and WMT14 En-De, with 153K and 4.5M bilingual sentence pairs in their respective training sets.Source: §A.1 · p. 12
  • IWSLT14 De-En uses a 7K validation split and concatenates dev2010, tst2010, tst2011, and tst2012 for testing; WMT14 En-De uses newstest2013 and newstest2014 for validation and testing.Source: §A.1 · p. 12
  • Both language pairs use BPE subword tokenization, shared source-target vocabularies and embeddings, and vocabularies extracted by Moses with default hyperparameters.Source: §A.1 · p. 12

A.2. Implementation Details

The experiments follow Transformer encoder-decoder architectures with task-specific dimensions and shared training settings. Evaluation uses tokenized BLEU, with an entropy upper bound that varies by sequence length.

A.2. Implementation Details

  • IWSLT14 De-En uses 6 encoder and decoder layers, 4 heads, and hidden and FFN dimensions of 512 and 1024; WMT14 En-De uses 8 heads and an FFN dimension of 2048.Both use dropout 0.1 on residual connections and attention weights.Source: §A.2 · p. 12
  • Training uses Adam with β1 = 0.9 and β2 = 0.98, label smoothing 0.1, a 4K-step warmup, batches of around 8192 tokens, and 10w training steps.The learning-rate schedule follows Vaswani et al. (2017).Source: §A.2 · p. 12
  • Tokenized BLEU is the evaluation metric; inference uses beam size 5 and task-specific length penalties, while the entropy upper bound is C = 0.7log(n).Length penalties are 1.1 for IWSLT14 De-En and 0.6 for WMT14 De-En; the entropy bound depends on sequence length n.Source: §A.2 · p. 12

B. Datasets of Document Translation

The document-translation evaluation uses Europarl7 En-De documents reconstructed at several length limits. The section also describes sparse activation and attention baselines for sentence-level translation.

B. Datasets of Document Translation

C. Baseline Details for Sentence-Level Translation

Loading 2302.06461v1…