PaperlayerView exact revision on arXiv

Source-linked AI summary

Bi-Directional Block Self-Attention for Fast and Memory-Efficient Sequence Modeling

Tao Shen, Tianyi Zhou, Guodong Long, Jing Jiang, Chengqi Zhang

arXiv:1804.00857v1cs.CLcs.AI

TL;DR

Existing sequence encoders must balance long-range and local context modeling against parallelism and memory, since self-attention’s alignment storage grows quadratically with sequence length. Bi-BloSAN addresses this by applying self-attention within blocks and then across block outputs, while adding feature-level attention and temporal masks. Across nine benchmark datasets, it achieves best or state-of-the-art performance with a better efficiency-memory trade-off than existing RNN/CNN/SAN models.

  • Problem

    Self-attention models local and long-range dependencies in parallel but requires memory that grows quadratically with sequence length, motivating a lower-memory RNN/CNN-free encoder.

  • Method

    Bi-BloSAN splits sequences into blocks, applies intra-block attention for local context and inter-block attention for long-range dependency, and uses feature-level attention with forward/backward masks.

  • Results

    Across nine benchmark datasets, Bi-BloSAN achieves best or state-of-the-art performance with a better efficiency-memory trade-off than existing RNN/CNN/SAN models.

  • Takeaways & Limitations

    Bi-BloSAN is much more time-efficient than RNN models, uses less memory than DiSAN, and outperforms CNN models and multi-head attention on prediction quality.

  • Takeaways & Limitations

    Additive attention usually achieves better empirical performance than multiplicative attention but is expensive in time cost and memory consumption.

Abstract

from arXiv · show

Recurrent neural networks (RNN), convolutional neural networks (CNN) and self-attention networks (SAN) are commonly used to produce context-aware representations. RNN can capture long-range dependency but is hard to parallelize and not time-efficient. CNN focuses on local dependency but does not perform well on some tasks. SAN can model both such dependencies via highly parallelizable computation, but memory requirement grows rapidly in line with sequence length. In this paper, we propose a model, called "bi-directional block self-attention network (Bi-BloSAN)", for RNN/CNN-free sequence encoding. It requires as little memory as RNN but with all the merits of SAN. Bi-BloSAN splits the entire sequence into blocks, and applies an intra-block SAN to each block for modeling local context, then applies an inter-block SAN to the outputs for all blocks to capture long-range dependency. Thus, each SAN only needs to process a short sequence, and only a small amount of memory is required. Additionally, we use feature-level attention to handle the variation of contexts around the same word, and use forward/backward masks to encode temporal order information. On nine benchmark datasets for different NLP tasks, Bi-BloSAN achieves or improves upon state-of-the-art accuracy, and shows better efficiency-memory trade-off than existing RNN/CNN/SAN.

1 INTRODUCTION

NLP sequence encoders trade off long-range context, parallelism, local modeling, and memory, motivating Bi-BloSAN as a block-structured self-attention alternative. Experiments across benchmark tasks report strong accuracy and efficiency-memory trade-offs.

  • RNNs capture long-range dependencies but are difficult to parallelize because of recurrent computation, whereas CNNs are parallelizable but struggle to relate distant positions efficiently.For hierarchical CNNs, the number of layers needed to connect arbitrary positions grows with their distance.
  • Self-attention flexibly models local and long-range dependencies with parallel computation, but storing all token-pair alignment scores requires memory that grows quadratically with sequence length.This motivates a self-attention encoder designed to retain SAN advantages while using memory closer to recurrent models.
  • Bi-BloSAN applies intra-block self-attention for local context and inter-block self-attention for long-range or global dependency modeling.Splitting sequences into blocks lets each attention operation process a shorter sequence than a single self-attention layer over the full input.
  • Bi-BloSA combines block-level attention with feature fusion, feature-level attention, and forward/backward masks to produce context-aware sequence representations.The masks encode temporal order, while feature-level attention handles variation in contexts around the same word.
  • Bi-BloSAN achieves the best accuracy with similar training time on SNLI while improving the efficiency-memory trade-off over existing RNN/CNN/SAN models.Figure 1 compares training time, training memory consumption, and test accuracy; the experiments cover nine benchmark datasets and multiple NLP tasks.

2 BACKGROUND

Sequence modeling begins with token embeddings, then uses attention mechanisms to build context-aware representations. Self-attention models token dependencies, while masking introduces directional order information.

  • 2.1 WORD EMBEDDING: Word embedding maps discrete tokens into low-dimensional real-valued vectors used as the basic sequence-modeling units.A pretrained embedding matrix transforms one-hot token vectors into representations.
  • 2.2 VANILLA ATTENTION AND MULTI-DIMENSIONAL ATTENTION: Vanilla attention scores each token against a query and converts the scores into probabilities indicating token importance.The output represents information selected according to these importance probabilities.
  • 2.2 VANILLA ATTENTION AND MULTI-DIMENSIONAL ATTENTION: Multi-dimensional attention assigns feature-specific scores, making it more expressive when a word’s meaning varies across contexts.Each feature has its own probability distribution over input tokens.
  • 2.3 TWO TYPES OF SELF-ATTENTION: Token2token self-attention replaces the query with another token and produces context-aware representations from pairwise token dependencies.Each output token receives a feature-level weighted combination of the input sequence.
  • 2.4 MASKED SELF-ATTENTION: Masked self-attention imposes one-way attention through a mask, enabling forward and backward variants that encode bidirectional order information.The mask modifies alignment scores, while fixed scaling and tanh define the masked compatibility function.

3 PROPOSED MODEL

Bi-BloSAN builds sequence representations with masked block self-attention, combining local intra-block context, long-range inter-block context, and gated feature fusion.

  • 3.2 BI-DIRECTIONAL BLOCK SELF-ATTENTION NETWORK FOR SEQUENCE ENCODING: Bi-BloSAN uses bi-directional block self-attention as its context-fusion module for sequence encoding.It combines forward and backward masked attention through Bi-BloSA.
  • 3.1 MASKED BLOCK SELF-ATTENTION: The model splits token embeddings into equal-length blocks and applies shared masked self-attention independently within each block to capture local dependencies.Padding may be added to the final block, and block length r is a hyper-parameter.
  • 3.1 MASKED BLOCK SELF-ATTENTION: Each block is summarized into a vector, and masked inter-block self-attention models long-range or global dependencies among those block representations.A shared source2token attention produces one local-context vector per block before inter-block attention.
  • 3.1 MASKED BLOCK SELF-ATTENTION: A gate dynamically combines the input and masked-attention output to merge local and global context features at block level.The resulting block-level global features are duplicated across tokens and fused with token embeddings and local features.
  • 3.1 MASKED BLOCK SELF-ATTENTION: The resulting mBloSA output consists of context-aware representations for all input tokens.Feature fusion combines the input sequence, local context features, and long-range/global context features.

4 EXPERIMENTS

Experiments across nine benchmark datasets compare Bi-BloSAN with RNN, CNN, and SAN-based encoders on accuracy, speed, and memory. Bi-BloSAN generally combines competitive or state-of-the-art prediction quality with faster computation and more favorable memory scaling.

  • 4.1 NATURAL LANGUAGE INFERENCE: On SNLI, Bi-BloSAN outperforms Bi-LSTM by 0.7%, Bi-GRU by 0.8%, Bi-SRU by 0.9%, multi-CNN by 2.5%, Hrchy-CNN by 1.8%, and multi-head attention by 1.5%.The comparison uses a 480D Bi-BloSAN encoder with a parameter count similar to baseline encoders.
  • 4.1 NATURAL LANGUAGE INFERENCE: Bi-BloSAN is 3∼4× faster than RNN models, matches CNNs and multi-head attention in speed, and uses much less training memory than DiSAN.On SNLI, it has similar GPU-memory use to RNN-based models and multi-head attention while substantially outperforming CNNs and multi-head attention in test accuracy.
  • 4.1 NATURAL LANGUAGE INFERENCE: The ablation shows that local and global context representations both contribute, while mBloSA raises test accuracy from 83.1% to 85.7%.Source2token self-attention improves test accuracy by 3.3% over vanilla attention.
  • 4 EXPERIMENTS: Across benchmark tasks, Bi-BloSAN achieves state-of-the-art or competitive prediction accuracy while improving efficiency relative to RNN/CNN/SAN baselines.It achieves state-of-the-art accuracy on the modified reading-comprehension task and state-of-the-art or best results across several semantic-relatedness and sentence-classification benchmarks.
  • 4.5 ANALYSES OF TIME COST AND MEMORY CONSUMPTION: As sequence length grows, Bi-BloSAN retains inference time similar to CNNs and multi-head attention while avoiding DiSAN's explosive memory growth.Its memory growth is nearly linear, and its block structure significantly reduces inference time compared with DiSAN.

5 CONCLUSIONS

Bi-BloSAN uses blockwise intra- and inter-block self-attention to model local and long-range dependencies while reducing memory demands. Experiments show strong performance and efficiency-memory trade-offs across NLP tasks.

  • Bi-BloSAN combines intra-block and inter-block self-attention to capture local and long-range context dependencies, respectively.
  • Forward and backward masks encode temporal order through asymmetric self-attentions.
  • Across nine benchmark datasets, Bi-BloSAN achieves best or state-of-the-art performance with a better efficiency-memory trade-off than existing RNN, CNN, and SAN models.

A THE SELECTION OF BLOCK LENGTH

The block length controls mBloSA memory consumption because attention cost depends on the squared lengths of its block sequences. The paper derives block-length selection for fixed-length data and normally distributed sentence lengths.

  • mBloSA memory is dominated by masked self-attentions whose cost is proportional to the square of sequence length.
  • For fixed sentence length n with m = n/r blocks, the paper minimizes memory by setting the gradient with respect to block length r to zero.
  • For normally distributed sentence lengths, block-length selection accounts for mini-batch size B and an upper bound on the expected maximum sentence length.

B NETWORK SETUP FOR MACHINE COMPREHENSION

The machine-comprehension setup converts SQuAD answer-span prediction into identifying the passage sentence containing the answer. Shared context-fusion and attention modules encode sentences and model sentence-level dependencies before scoring candidates.

  • The modified SQuAD task predicts which passage sentence contains the correct answer rather than the answer span.
  • Each passage sentence is encoded with a shared context-fusion layer and source2token self-attention into a sentence vector.
  • Sentence and question representations are combined as [u_k; q; u_k−q; u_k ⊙q], then processed for sentence-level dependencies and answer scoring.
  • The context-fusion candidates include Bi-LSTM, Bi-GRU, Bi-SRU, multi-CNN, multi-head attention, and Bi-BloSA.

C LOSS OF REGRESSION PROBLEM

The regression setup represents similarity degrees as probability distributions and trains predictions against those targets using KL-divergence.

  • The regression model predicts similarity degree from a feedforward network’s probability distribution using β = [1, 2, . . . , K].
  • Ground-truth similarity y is mapped to a target distribution p whose expectation satisfies y = β^T p.
  • KL-divergence between target p and predicted distribution p-hat is used as the training loss.

D RELATED WORKS

The paper situates Bi-BloSA among structured and self-attention mechanisms, distinguishing it from hierarchical attention by its token-level context encoding and masked two-stage self-attention.

  • Structured attention mechanisms model dependencies between tokens and relate to self-alignment and multi-head attention with one head.
  • Bi-BloSA applies masked self-attention twice: intra-block attention for local context and inter-block attention for global context.It also uses feature fusion gates and positional masks to combine contexts and encode temporal order.
  • Unlike hierarchical attention, Bi-BloSA learns context-aware representations for each token rather than document embeddings.Its input is a sentence, whereas hierarchical attention processes documents composed of multiple sentences.
  • Bi-BloSA and hierarchical attention share a two-layer bottom-to-top attention structure, but differ in attention type and processing granularity.Hierarchical attention uses vanilla token-level and sentence-level attention, while Bi-BloSA uses masked self-attention across blocks.
Loading 1804.00857v1…