PaperlayerView exact revision on arXiv

Source-linked AI summary

Adaptively Sparse Transformers

Gonçalo M. Correia, Vlad Niculae, André F. T. Martins

arXiv:1909.00015v2cs.CLstat.ML

TL;DR

Softmax attention gives every context word non-zero weight, limiting sparsity in Transformer heads. The paper replaces it with α-entmax and learns α per head, finding more interpretable and diverse attention with no accuracy cost. Its analyses show distinct sparsity preferences and head specializations, while decoder self-attention may favor denser distributions because it sees only past tokens.

  • Problem

    Softmax attention assigns non-zero weight to all context words, motivating sparse attention for more selective and interpretable attention distributions.

  • Method

    The paper replaces Transformer softmax attention with α-entmax and derives gradients that allow α, controlling sparsity, to be learned automatically for each head.

  • Results

    The adaptively sparse Transformer improves interpretability and head diversity without sacrificing accuracy, while different heads learn different sparsity behaviors.

  • Takeaways & Limitations

    Sparsity helps uncover specialized behaviors, including near-deterministic positional heads and a subword-joining head.

  • Takeaways & Limitations

    Decoder self-attention may prefer denser distributions because its autoregressive access to past tokens provides fewer tokens to attend to.

Abstract

from arXiv · show

Attention mechanisms have become ubiquitous in NLP. Recent architectures, notably the Transformer, learn powerful context-aware word representations through layered, multi-headed attention. The multiple heads learn diverse types of word relationships. However, with standard softmax attention, all attention heads are dense, assigning a non-zero weight to all context words. In this work, we introduce the adaptively sparse Transformer, wherein attention heads have flexible, context-dependent sparsity patterns. This sparsity is accomplished by replacing softmax with $α$-entmax: a differentiable generalization of softmax that allows low-scoring words to receive precisely zero weight. Moreover, we derive a method to automatically learn the $α$ parameter -- which controls the shape and sparsity of $α$-entmax -- allowing attention heads to choose between focused or spread-out behavior. Our adaptively sparse Transformer improves interpretability and head diversity when compared to softmax Transformers on machine translation datasets. Findings of the quantitative and qualitative analysis of our approach include that heads in different layers learn different sparsity preferences and tend to be more diverse in their attention distributions than softmax Transformers. Furthermore, at no cost in accuracy, sparsity in attention heads helps to uncover different head specializations.

1 Introduction

Transformers use multi-head attention to model diverse token relationships, but softmax assigns non-zero weight to every context word. The paper introduces adaptively sparse attention with learnable sparsity patterns to improve interpretability and head diversity without sacrificing accuracy.

  • Transformer attention heads learn diverse token relationships, but softmax assigns non-zero weight to every context word.
  • The proposed Transformer attends to sparse, potentially non-contiguous sets of words and adapts each head’s attended support.
  • The model introduces sparse attention into Transformers, easing interpretability and producing slight accuracy gains.
  • Adaptive sparse attention lets each head vary continuously between dense softmax behavior and sparse, piecewise-linear sparsemax behavior.
  • The paper analyzes sharper existing and novel attention-head behaviors revealed by sparsity and adaptivity.

2 Background

The Transformer propagates information through hierarchical multi-head attention, while α-entmax provides a flexible normalization that can produce sparse attention. The paper derives gradients for learning α automatically, enabling different heads to adopt different sparsity behaviors.

  • 2.1 The Transformer: In NMT, the Transformer maps input sequences to output sequences using hierarchical multi-head attention for dynamic, context-dependent information flow.
  • 2.1 The Transformer: Scaled dot-product attention uses keys to score items, softmax to normalize weights, and values to form weighted representations.
  • 2.1 The Transformer: Multi-head attention applies attention in parallel with distinct learned linear transformations, allowing heads to specialize.
  • 2.2 Sparse Attention: α-entmax is a Tsallis-entropy family containing Shannon and Gini entropy cases at α = 1 and α = 2.
  • 2.2 Sparse Attention: At α = 1, α-entmax recovers softmax; for α > 1 it permits sparse solutions, with α = 2 recovering piecewise-linear sparsemax.
  • 2.2 Sparse Attention: The authors derive a Jacobian with respect to α, enabling automatic gradient-based learning of attention shape and sparsity.

3 Adaptively Sparse Transformers with α-entmax

The paper replaces softmax in Transformer attention heads with α-entmax to produce sparse, flexible attention patterns. It also learns a separate α for each head, enabling heads to vary continuously between dense and sparse behavior.

  • The proposed Transformer replaces the attention row normalization with α-entmax.
  • Different α per head: Each attention head receives an adaptive α value because heads can capture different linguistic phenomena and either isolate words or spread attention across phrases.
  • Derivatives w.r.t. α: A closed-form Jacobian of entmax with respect to α enables adaptive α training through gradient methods.The paper describes this derivative as a key contribution and uses it to train adaptively sparse Transformers.

4 Experiments

The experiments evaluate fixed and adaptive entmax Transformers against the standard softmax Transformer on four machine-translation datasets. Replacing softmax with entmax does not hurt performance and tends to produce slightly higher BLEU while improving analysis potential through sparsity.

  • Experimental setup: The evaluation covers four machine-translation tasks using the standard softmax Transformer as the natural baseline.The datasets are IWSLT 2017 German–English, KFTT Japanese–English, WMT 2016 Romanian–English, and WMT 2014 English–German.
  • Model variants: The fixed 1.5-entmax model uses sparse entmax attention with α = 1.5 for every head.The paper presents this as a Transformer model distinct from earlier RNN-based use of 1.5-entmax.
  • Model variants: The adaptive α-entmax model learns a different α for each attention head and layer.It adds scalar α parameters for the encoder self-attention, context attention, and decoder self-attention mechanisms.
  • Experimental setup: The models are trained on BPE-preprocessed datasets using joint segmentations with 32k merge operations.
  • Experimental setup: The proposed methods incur relatively small computational overhead compared with standard softmax.The passage reports training-throughput comparisons for α-entmax and 1.5-entmax but does not provide the complete values.
  • Results: Entmax does not hurt test-set tokenized BLEU on any dataset and tends to produce slightly higher BLEU than softmax.The reported benefit is accompanied by greater potential for analyzing the learned attention heads.

5 Analysis

Analysis shows that learned α values produce layer- and head-specific sparsity, greater attention diversity, and interpretable specializations without sacrificing translation performance.

  • 5.1 High-Level Statistics: During training, most randomly initialized α values first decrease, then some heads become sparser after about one thousand steps.This suggests softmax-like behavior may be preferred while the model is uncertain, followed by specialization as confidence increases.
  • 5.1 High-Level Statistics: At convergence, encoder self-attention develops sparse and dense α modes, whereas decoder self-attention and context attention remain unimodal.The sparse encoder mode approaches α →2, while the dense mode lies between softmax and 1.5-entmax.
  • 5.2 Identifying Head Specializations: Learned α combines fully dense and sparse attention, with encoder self-attention tending toward extreme sparsity and the other modules showing more uniform densities.Compared with fixed α = 1.5, learned α also creates clearer differences among modules and more varied head behavior across layers.
  • 5.2 Identifying Head Specializations: Sparse entmax models show greater head diversity than the softmax baseline, with diversity peaking in middle encoder and context-attention layers.Diversity is measured using Jensen-Shannon divergence, where higher values indicate greater disagreement among heads.
  • 5.2 Identifying Head Specializations: Adaptive sparsity strengthens positional-head specialization: four adaptive heads attend to the previous token with median confidence 95.9%, versus 93.5% for softmax and 94.4% for 1.5-entmax.The adaptive model includes both a dense lower-confidence head with α = 1.18 and a sparse higher-confidence head with α = 1.91.
  • 5.2 Identifying Head Specializations: Sparse models reveal a BPE-merging head that spreads attention over neighboring tokens within BPE clusters or hyphenated words.Two heads in each entmax model exceed 80% BPE-merging score, compared with none in the softmax model; α-entmax reaches 88.5% and 89.8%.
  • 5.2 Identifying Head Specializations: An adaptive interrogation head assigns 99.5% average attention to the question mark for interrogative words, compared with 98.5% for softmax and 97.0% for 1.5-entmax.The adaptive head chooses α close to 1, making it closer to softmax.
  • 5.2 Identifying Head Specializations: The same head can exhibit different sparsity on similarly long sentences when phrase confidence differs.The phrase “a sexually transmitted disease” receives higher confidence and therefore more sparsity in the example.

6 Related Work

The paper distinguishes its adaptive sparse attention from prior fixed or span-limited mechanisms, while extending entmax with dynamically learned sparsity and analyzing its interpretability.

  • Sparse attention: Unlike prior entmax work using one fixed α, this model automatically adjusts a continuous α parameter to vary curvature and sparsity dynamically.The paper presents this as the first work allowing such attention mappings to adapt dynamically and applying sparse attention in a Transformer.
  • Fixed sparsity patterns: Earlier sparse mechanisms limit attention to contiguous spans, whereas this model can attend to sparse, non-contiguous sets of words.The paper contrasts its approach with fixed sparsity patterns and adaptive attention spans.
  • Transformer interpretability: The paper uses sparse attention to make head roles easier to inspect and reports both clearer known behaviors and novel specializations.Its interpretability analysis examines attention-head behavior in Transformer models.

7 Conclusion and Future Work

The paper presents adaptively sparse attention and analyzes its potential for translation accuracy and interpretability. It finds that adaptive sparsity supports more specialized attention behavior while avoiding costly per-head hyperparameter searches.

  • The authors introduce a strategy for adaptively sparse attention, specifically for adaptively sparse Transformers.
  • Sparse attention mappings provide potential benefits for both translation accuracy and model interpretability.
  • Gradient-based optimization learns attention-head adaptivity without costly per-head hyperparameter searches.
  • Adaptive sparsity allows attention heads to specialize more and with higher confidence.
  • Automatically learned behaviors, including near-deterministic positional heads and a subword-joining head, may inspire static Transformer variants.

A High-Level Statistics Analysis of Other Language Pairs

This section presents high-level analyses of learned shape parameters, head densities, and layer-wise divergence, alongside the mathematical relationships among softmax, sparsemax, and α-entmax. The analyses characterize how α controls sparsity and curvature.

  • Histograms summarize learned α values and head densities, while layer-wise plots track Jensen-Shannon divergence and head density.
  • Sparsemax is the Euclidean projection onto the simplex and has piecewise-linear coordinates.
  • Softmax is characterized as the unique solution of an entropy-regularized optimization problem, with strictly positive probabilities from its exponential form.

C Jacobian of α-entmax w.r.t. the shape parameter α: Proof of Proposition 1

This section derives the Jacobian of α-entmax with respect to its shape parameter α. It begins from the entmax transformation and its associated Shannon-entropy formulation.

  • The entmax transformation is introduced as the starting point for the derivation.
  • The Shannon entropy term appears in the optimization formulation underlying the transformation.
  • The section derives the Jacobian of entmax with respect to the scalar parameter α.

C.1 General case of α > 1

The general α > 1 case derives the α-entmax Jacobian using the solution’s threshold characterization and implicit differentiation. The derivation handles support-restricted coordinates, indeterminate expressions, and nondifferentiable points.

  • For general α, the α-entmax solution is expressed coordinate-wise through the KKT conditions of its optimization problem.
  • The scalar threshold τ⋆ is defined implicitly so that the solution p⋆ normalizes to 1.
  • For general α, the defining equation lacks a closed-form solution, making Jacobian computation non-trivial.
  • Implicit differentiation provides a way to obtain the Jacobian despite the lack of a closed-form solution.
  • The Jacobian exists almost everywhere, with a generalized Jacobian available at certain nondifferentiable (α, z) pairs.
  • Because increasing α keeps sparse coordinates sparse, the derivation focuses on coordinates within the support of p⋆.
  • The remaining derivation resolves indeterminate expressions using L’Hôpital’s rule and combines the resulting terms into the Jacobian.
  • The section concludes with an explicit expression for the Jacobian of α-entmax with respect to α, including the α = 1 case.
Loading 1909.00015v2…