Feed-Forward Networks with Attention Can Solve Some Long-Term Memory Problems

In plain words

Models that read a sequence, such as text or audio, often go through it one step at a time, remembering what came before. This makes very long sequences slow to process and hard to remember over long distances. Attention is a way for a model to look back at any earlier part of the input directly.

This paper asks whether attention can be made much simpler. Each step of the input gets a single importance score, computed from that step alone. The scores are turned into weights that add up to one, and the model keeps a weighted average of the steps as one summary vector. Since no step depends on the previous one, every step can be processed at the same time.

The authors test this on two synthetic memory tasks, addition and multiplication, with sequences longer and more varied in length than earlier published results on these tasks. They also compare the weighted version with a plain average over time, which is the simplest way to combine steps. The weighted version solves the tasks at every length tested; the plain average is slower on some tasks and fails on multiplication at the longest lengths and when lengths vary. The paper is clear that these tasks cannot test word order, so the method is not shown to help where order matters.

TL;DR

• Feed-forward attention: a per-step scalar score a(h_t), a softmax over time, and a weighted sum of the step states gives one vector per sequence §1.2 eq. 1.
• The model has no recurrence, so each step is computed independently and in parallel §1.2.
• On the addition and multiplication toy tasks, it reaches perfect accuracy within the tested lengths T0 = 50 to 10000 (author-reported) tab. 1.
• With one model trained on lengths from 50 to 10000, it reaches 99.9% (addition) and 99.4% (multiplication) test accuracy (author-reported) §2.3.
• The unweighted mean is a much weaker baseline on the same tasks (77.4% and 55.5% on variable lengths) §2.3.

Problem

Recurrent networks are trained by backpropagation through time, which suffers from vanishing and exploding gradients. In practice they are used only where dependencies span at most hundreds of steps §1. Sequential evaluation also prevents full parallelism on long inputs §1.

Attention lets a model relate distant states directly. The paper asks whether attention can work in a feed-forward model with no recurrence, so that variable-length sequences are handled and computation can be parallelised §1.2. The experiments use the synthetic addition and multiplication long-term memory problems, whose exact specification is left to earlier work §2. These tasks are commutative, so a model that ignores order can solve them in principle §2. The authors state that the model cannot solve tasks where temporal order matters §2.4.

Method

Overview. Each input step is mapped to a hidden state with no dependence on other steps. Feed-forward attention then scores each state with a learned scalar function, normalises the scores over time, and averages the states with those weights. The result is a fixed-length vector for each sequence. The rest of the network is a small feed-forward stack on that vector §2.1.

Steps.

1. Each input step x_t is a vector of 2 values §2.1. Its state is h_t = LReLU(W_xh x_t + b_xh) with D = 100 §2.1. LReLU(x) = max(x, 0.01x) §2.1.
2. The score of each state is e_t = a(h_t), where a(h) = tanh(W_hc h + b_hc) with W_hc in R^{1×D} and b_hc a scalar §2.1 eq. 1.
3. The weights α_t are a softmax of the scores over all T steps, and the context is c = Σ_t α_t h_t eq. 1 §1.2.
4. The context goes through s = LReLU(W_cs c + b), then y = LReLU(W_sy s + b), where y is the single output value §2.1.
5. The baseline replaces step 3 with the unweighted mean c = (1/T) Σ_t h_t eq. 2 §1.2. Everything else is the same §2.1.

The general attention formulation with a dependence on the previous decoder state is given in §1.1. It is not used here.

# Feed-forward attention pooling  [eq. 1] [§2.1]
# Input:  inputs (B, T, 2), two values per step
# Output: context (B, D), one vector per sequence

def feed_forward_attention(inputs):
    # Each state depends only on its own step  [§2.1]
    states = leaky_relu(inputs @ W_xh.T + b_xh)   # (B, T, D)
    # Scalar score per step, bounded by tanh  [§2.1] [eq. 1]
    scores = tanh(states @ W_hc.T + b_hc)         # (B, T, 1)
    # Normalise the scores over time  [eq. 1]
    weights = exp(scores) / sum(exp(scores), axis=1)
    # Weighted average of the states  [eq. 1]
    return sum(weights * states, axis=1)          # (B, D)

# Training loop for the toy problems  [§2.1]
# Input:  learning_rate from {3e-4, 1e-3, 3e-3, 1e-2}  [§2.2]
# Output: accuracy history; stop at 100% or after 100 epochs

def train(learning_rate):
    # Adam with beta1 = 0.9 and beta2 = 0.999  [§2.1]
    optimizer = adam(learning_rate, beta1=0.9, beta2=0.999)
    # One epoch is 1000 parameter updates  [§2.1]
    for update in range(100 * 1000):
        # Minibatch of 100 sequences, squared error on the output  [§2.1]
        inputs, targets = sample_minibatch(batch_size=100)
        optimizer.step(squared_error(model(inputs), targets))
        if (update + 1) % 1000 == 0:
            # Accuracy on 1000 held-out sequences  [§2.1]
            accuracy = mean(abs(model(test_inputs) - test_targets) < 0.04)
            if accuracy == 1.0:
                # Stop when the test set is perfect  [§2.1]
                return accuracy
    return accuracy

Hyperparameters

Results

Author-reported. Setup: synthetic addition and multiplication tasks §2; one model with D = 100 §2.1; baseline is the unweighted mean of the same model §2.2. Timing hardware is an NVIDIA GTX 980 Ti §2.2.

Fixed-length experiment tab. 1. Each cell is the number of epochs needed to reach perfect accuracy. Grey cells in the source are accuracy after 100 epochs instead, given as a percentage.

Source: tab. 1 and §2.2. The authors conclude that the attention model solves these tasks quickly at every length tested §2.2.

Speed §2.2. One epoch of 100,000 sequences at T0 = 10000 took 254 s with the feed-forward model and 917 s with a single-layer vanilla RNN of hidden size 100 on the same GPU. The paper says both models have comparable parameter counts.

Parameter counts §2.2. 10,602 with attention and 10,501 with the unweighted mean. The difference is under 1%.

Variable-length experiment §2.3. One model trained on lengths drawn uniformly from 50 to 10000:

The paper does not report baseline numbers from published RNN results. It says earlier RNNs were trained to lengths of a few hundred to a few thousand §2.2 and that its claim is about lengths "longer and more widely varying" than the best published results [abstract].

Paper vs code

Code: official repository https://github.com/craffel/ff-attention, commit e8fce3e66149d18d8766dfe703100e407c34aaea. The paper's footnote links it §2.4. The repository has no LICENSE file, so its code is cited here and not quoted.

• Weight initialisation differs. The paper says all weight matrices are Gaussian with std 1/sqrt(N) §2.1. The attention weight uses that scheme fixed_length_experiments.py. The dense layers use lasagne.init.HeNormal() fixed_length_experiments.py fixed_length_experiments.py fixed_length_experiments.py. Lasagne defines the scale of HeNormal; that definition is not in this repository, so the std is not determined from these sources.
• Variable-length lengths are per minibatch, not per sequence. The paper says minibatches of sequences have lengths chosen uniformly at random §2.3. The code draws one length for the whole minibatch, and every sequence in that batch uses it variable_length_experiments.py. The per-sequence reading in the paper's wording is therefore not what the code does.
• Early stop beyond the paper's rule. The paper stops at 100% accuracy or 100 epochs §2.1. The code also stops a run once its batch count passes the best batch count found so far for the same task and aggregation fixed_length_experiments.py variable_length_experiments.py. The fixed-length script keys this by sequence length too; the variable-length script does not. The two scripts also use >= and > respectively. Reported epochs for runs that were pruned are therefore caps, not the epochs each run needed.
• Units in the output. The paper reports epochs of 1000 updates §2.1 tab. 1. The code writes the 0-based batch index of the final update, not epochs fixed_length_experiments.py. Epochs can be recovered by dividing by 1000, which is inferred from the code and the paper's definition.
• Evaluation of the mask. The data generator returns a mask. Neither script passes it to the network fixed_length_experiments.py variable_length_experiments.py. The paper's model sums over all T steps eq. 1, so the mask matters only if padding is present. Lengths in the fixed-length run are drawn from [T0, 1.1·T0] and then padded to a common size §2.2. Whether padded steps enter the sum depends on the data generator (see Pitfalls).
• Input dimension and padded length are set by the external lstm_problems package, which is not in this repository. Lasagne's input shape is taken from one dummy batch fixed_length_experiments.py, so the input's sequence dimension is fixed to that batch's length, while training batches vary in length. How the code handles this at run time is not verified here.
• Adam epsilon and leaky slope are not set in the code. The paper does not state epsilon §2.1. The code uses library defaults fixed_length_experiments.py fixed_length_experiments.py. The leaky slope of 0.01 matches the paper, but the match comes from the library default, not from the code.
• Parameter count consistency. The paper's 10,602 and 10,501 §2.2 match the architecture with 2 input features and D = 100. Input layer 300, attention 101, intermediate layer 10,100, output 101 gives 10,602; without attention, 10,501. This check is derived here, not stated in the paper or code.
• No discrepancies found for: hidden size D = 100, batch size 100, learning rate grid, the 0.04 accuracy threshold, the tanh score function, the softmax over time, and the number of test sequences.

Pitfalls

• Order is discarded. The pooled vector does not depend on the order of the steps eq. 1 §2.4. The paper's own counter-example is the two-symbol order task, where X, Y and Y, X cannot be told apart §2.4. The toy tasks are commutative §2, so they can be solved without order information.
• Masking is the implementer's responsibility. The equations sum over T steps. If you pad sequences, mask padded positions in both the softmax and the mean. A zero input does not give a zero state: h = LReLU(b) is nonzero when the bias is nonzero. The reference code does not pass its mask to the network, and the external generator's padding behaviour is not verified here.
• Per-minibatch lengths in the variable-length run. The paper's wording does not settle whether lengths vary per sequence or per minibatch. The code uses one length per minibatch variable_length_experiments.py, so reproduce that unless the authors confirm otherwise.
• The stopping rule of the code is not the paper's rule. Runs that were pruned report the accuracy at the time of pruning, not 100% fixed_length_experiments.py. Compare Table 1 with that in mind tab. 1.
• Initialisation. Use the paper's stated scheme (std 1/sqrt(N) §2.1) if you want to match the paper's text. The code uses HeNormal instead, and its scale is set by Lasagne fixed_length_experiments.py.
• Softmax stability. The reference code applies exp to a tanh-bounded score, so it cannot overflow layers.py. If you change the score nonlinearity, subtract the maximum score before exp.
• Single-run results. The paper reports no seeds, repeats or variance §2.2. The scripts do not set a random seed fixed_length_experiments.py. Treat single-table cells as indicative.
• Epsilon and other defaults. Adam epsilon, the leaky slope, dropout, weight decay, gradient clipping and any learning-rate schedule are not stated in the paper. The code relies on library defaults for the first two and has none for the others.
• Metric is thresholded. Accuracy is the fraction of sequences with absolute error below 0.04 §2, and training stops on accuracy, not on loss §2.1.
• The code is Python 2 and needs Theano and Lasagne fixed_length_experiments.py README.md. The data generator lstm_problems is a separate repository, and this card's claims about its padding and feature layout come from a summary of that repository, not from a pinned copy. Verify before relying on them.
• Version. Section and equation numbers follow the LaTeX source that wikillm fetch downloaded. The arXiv metadata lists version v5 (updated 2016-09-20).

Lineage

• Builds on: attention as defined for sequence models by Bahdanau et al. §1.1, arXiv:1409.0473. The paper's change is to drop the dependence on the previous decoder state and apply the score per step §1.2.
• Long-term memory tasks: addition and multiplication are the synthetic tasks of Hochreiter and Schmidhuber, as the paper cites them §2.
• Compared against: the unweighted mean over time, the same model with a fixed average eq. 2 §2.2; a single-layer vanilla RNN of hidden size 100, for speed only §2.2.
• Related work the paper names: Sønderby et al. compared a recurrent network with and without this simplified attention on protein sequence analysis §1.2; mean pooling over time appears in the audio and text work it cites §1.2.
• Superseded by: not stated in this wiki.