Bahdanau Attention

:label:sec_seq2seq_attention

We studied the machine translation problem in :numref:sec_seq2seq, where we designed an encoder-decoder architecture based on two RNNs for sequence to sequence learning. Specifically, the RNN encoder transforms a variable-length sequence into a fixed-shape context variable, then the RNN decoder generates the output (target) sequence token by token based on the generated tokens and the context variable. However, even though not all the input (source) tokens are useful for decoding a certain token, the same context variable that encodes the entire input sequence is still used at each decoding step.

In a separate but related challenge of handwriting generation for a given text sequence, Graves designed a differentiable attention model to align text characters with the much longer pen trace, where the alignment moves only in one direction :cite:Graves.2013. Inspired by the idea of learning to align, Bahdanau et al. proposed a differentiable attention model without the severe unidirectional alignment limitation :cite:Bahdanau.Cho.Bengio.2014. When predicting a token, if not all the input tokens are relevant, the model aligns (or attends) only to parts of the input sequence that are relevant to the current prediction. This is achieved by treating the context variable as an output of attention pooling.

Model

When describing Bahdanau attention for the RNN encoder-decoder below, we will follow the same notation in :numref:sec_seq2seq. The new attention-based model is the same as that in :numref:sec_seq2seq except that the context variable $\mathbf{c}$ in :eqref:eq_seq2seq_s_t is replaced by $\mathbf{c}_{t’}$ at any decoding time step $t’$. Suppose that there are $T$ tokens in the input sequence, the context variable at the decoding time step $t’$ is the output of attention pooling:

\mathbf{c}{t’} = \sum{t=1}^T \alpha(\mathbf{s}_{t’ - 1}, \mathbf{h}_t) \mathbf{h}_t,

where the decoder hidden state $\mathbf{s}_{t’ - 1}$ at time step $t’ - 1$ is the query, and the encoder hidden states $\mathbf{h}_t$ are both the keys and values, and the attention weight $\alpha$ is computed as in :eqref:eq_attn-scoring-alpha using the additive attention scoring function defined by :eqref:eq_additive-attn.

Slightly different from the vanilla RNN encoder-decoder architecture in :numref:fig_seq2seq_details, the same architecture with Bahdanau attention is depicted in :numref:fig_s2s_attention_details.

Layers in an RNN encoder-decoder model with Bahdanau attention. :label:fig_s2s_attention_details

```{.python .input} from d2l import mxnet as d2l from mxnet import np, npx from mxnet.gluon import rnn, nn npx.set_np()

  2. ```{.python .input}
  3. #@tab pytorch
  4. from d2l import torch as d2l
  5. import torch
  6. from torch import nn

Defining the Decoder with Attention

To implement the RNN encoder-decoder with Bahdanau attention, we only need to redefine the decoder. To visualize the learned attention weights more conveniently, the following AttentionDecoder class defines the base interface for decoders with attention mechanisms.

```{.python .input}

@tab all

@save

class AttentionDecoder(d2l.Decoder): “””The base attention-based decoder interface.””” def init(self, kwargs): super(AttentionDecoder, self).init(kwargs)

  1. @property
  2. def attention_weights(self):
  3.     raise NotImplementedError
  2. Now let us implement
  3. the RNN decoder with Bahdanau attention
  4. in the following `Seq2SeqAttentionDecoder` class.
  5. The state of the decoder
  6. is initialized with 
  7. i) the encoder final-layer hidden states at all the time steps (as keys and values of the attention);
  8. ii) the encoder all-layer hidden state at the final time step (to initialize the hidden state of the decoder);
  9. and iii) the encoder valid length (to exclude the padding tokens in attention pooling).
  10. At each decoding time step,
  11. the decoder final-layer hidden state at the previous time step is used as the query of the attention.
  12. As a result, both the attention output
  13. and the input embedding are concatenated
  14. as the input of the RNN decoder.
  16. ```{.python .input}
  17. class Seq2SeqAttentionDecoder(AttentionDecoder):
  18.     def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
  19.                  dropout=0, **kwargs):
  20.         super(Seq2SeqAttentionDecoder, self).__init__(**kwargs)
  21.         self.attention = d2l.AdditiveAttention(num_hiddens, dropout)
  22.         self.embedding = nn.Embedding(vocab_size, embed_size)
  23.         self.rnn = rnn.GRU(num_hiddens, num_layers, dropout=dropout)
  24.         self.dense = nn.Dense(vocab_size, flatten=False)
  26.     def init_state(self, enc_outputs, enc_valid_lens, *args):
  27.         # Shape of `outputs`: (`num_steps`, `batch_size`, `num_hiddens`).
  28.         # Shape of `hidden_state[0]`: (`num_layers`, `batch_size`,
  29.         # `num_hiddens`)
  30.         outputs, hidden_state = enc_outputs
  31.         return (outputs.swapaxes(0, 1), hidden_state, enc_valid_lens)
  33.     def forward(self, X, state):
  34.         # Shape of `enc_outputs`: (`batch_size`, `num_steps`, `num_hiddens`).
  35.         # Shape of `hidden_state[0]`: (`num_layers`, `batch_size`,
  36.         # `num_hiddens`)
  37.         enc_outputs, hidden_state, enc_valid_lens = state
  38.         # Shape of the output `X`: (`num_steps`, `batch_size`, `embed_size`)
  39.         X = self.embedding(X).swapaxes(0, 1)
  40.         outputs, self._attention_weights = [], []
  41.         for x in X:
  42.             # Shape of `query`: (`batch_size`, 1, `num_hiddens`)
  43.             query = np.expand_dims(hidden_state[0][-1], axis=1)
  44.             # Shape of `context`: (`batch_size`, 1, `num_hiddens`)
  45.             context = self.attention(
  46.                 query, enc_outputs, enc_outputs, enc_valid_lens)
  47.             # Concatenate on the feature dimension
  48.             x = np.concatenate((context, np.expand_dims(x, axis=1)), axis=-1)
  49.             # Reshape `x` as (1, `batch_size`, `embed_size` + `num_hiddens`)
  50.             out, hidden_state = self.rnn(x.swapaxes(0, 1), hidden_state)
  51.             outputs.append(out)
  52.             self._attention_weights.append(self.attention.attention_weights)
  53.         # After fully-connected layer transformation, shape of `outputs`:
  54.         # (`num_steps`, `batch_size`, `vocab_size`)
  55.         outputs = self.dense(np.concatenate(outputs, axis=0))
  56.         return outputs.swapaxes(0, 1), [enc_outputs, hidden_state,
  57.                                         enc_valid_lens]
  59.     @property
  60.     def attention_weights(self):
  61.         return self._attention_weights

```{.python .input}

@tab pytorch

class Seq2SeqAttentionDecoder(AttentionDecoder): def init(self, vocabsize, embedsize, num_hiddens, num_layers, dropout=0, **kwargs): super(Seq2SeqAttentionDecoder, self).__init(**kwargs) self.attention = d2l.AdditiveAttention( num_hiddens, num_hiddens, num_hiddens, dropout) self.embedding = nn.Embedding(vocab_size, embed_size) self.rnn = nn.GRU( embed_size + num_hiddens, num_hiddens, num_layers, dropout=dropout) self.dense = nn.Linear(num_hiddens, vocab_size)

  1. def init_state(self, enc_outputs, enc_valid_lens, *args):
  2.     # Shape of `outputs`: (`num_steps`, `batch_size`, `num_hiddens`).
  3.     # Shape of `hidden_state[0]`: (`num_layers`, `batch_size`,
  4.     # `num_hiddens`)
  5.     outputs, hidden_state = enc_outputs
  6.     return (outputs.permute(1, 0, 2), hidden_state, enc_valid_lens)
  8. def forward(self, X, state):
  9.     # Shape of `enc_outputs`: (`batch_size`, `num_steps`, `num_hiddens`).
  10.     # Shape of `hidden_state[0]`: (`num_layers`, `batch_size`,
  11.     # `num_hiddens`)
  12.     enc_outputs, hidden_state, enc_valid_lens = state
  13.     # Shape of the output `X`: (`num_steps`, `batch_size`, `embed_size`)
  14.     X = self.embedding(X).permute(1, 0, 2)
  15.     outputs, self._attention_weights = [], []
  16.     for x in X:
  17.         # Shape of `query`: (`batch_size`, 1, `num_hiddens`)
  18.         query = torch.unsqueeze(hidden_state[-1], dim=1)
  19.         # Shape of `context`: (`batch_size`, 1, `num_hiddens`)
  20.         context = self.attention(
  21.             query, enc_outputs, enc_outputs, enc_valid_lens)
  22.         # Concatenate on the feature dimension
  23.         x = torch.cat((context, torch.unsqueeze(x, dim=1)), dim=-1)
  24.         # Reshape `x` as (1, `batch_size`, `embed_size` + `num_hiddens`)
  25.         out, hidden_state = self.rnn(x.permute(1, 0, 2), hidden_state)
  26.         outputs.append(out)
  27.         self._attention_weights.append(self.attention.attention_weights)
  28.     # After fully-connected layer transformation, shape of `outputs`:
  29.     # (`num_steps`, `batch_size`, `vocab_size`)
  30.     outputs = self.dense(torch.cat(outputs, dim=0))
  31.     return outputs.permute(1, 0, 2), [enc_outputs, hidden_state,
  32.                                       enc_valid_lens]
  34. @property
  35. def attention_weights(self):
  36.     return self._attention_weights
  2. In the following, we test the implemented 
  3. decoder with Bahdanau attention
  4. using a minibatch of 4 sequence inputs
  5. of 7 time steps.
  7. ```{.python .input}
  8. encoder = d2l.Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16,
  9.                              num_layers=2)
  10. encoder.initialize()
  11. decoder = Seq2SeqAttentionDecoder(vocab_size=10, embed_size=8, num_hiddens=16,
  12.                                   num_layers=2)
  13. decoder.initialize()
  14. X = d2l.zeros((4, 7))  # (`batch_size`, `num_steps`)
  15. state = decoder.init_state(encoder(X), None)
  16. output, state = decoder(X, state)
  17. output.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape

```{.python .input}

@tab pytorch

encoder = d2l.Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16, num_layers=2) encoder.eval() decoder = Seq2SeqAttentionDecoder(vocab_size=10, embed_size=8, num_hiddens=16, num_layers=2) decoder.eval() X = d2l.zeros((4, 7), dtype=torch.long) # (batch_size, num_steps) state = decoder.init_state(encoder(X), None) output, state = decoder(X, state) output.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape

  2. ## Training
  5. Similar to :numref:`sec_seq2seq_training`,
  6. here we specify hyperparemeters,
  7. instantiate
  8. an encoder and a decoder with Bahdanau attention,
  9. and train this model for machine translation.
  10. Due to the newly added attention mechanism,
  11. this training is much slower than
  12. that in :numref:`sec_seq2seq_training` without attention mechanisms.
  14. ```{.python .input}
  15. #@tab all
  16. embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1
  17. batch_size, num_steps = 64, 10
  18. lr, num_epochs, device = 0.005, 250, d2l.try_gpu()
  20. train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)
  21. encoder = d2l.Seq2SeqEncoder(
  22.     len(src_vocab), embed_size, num_hiddens, num_layers, dropout)
  23. decoder = Seq2SeqAttentionDecoder(
  24.     len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)
  25. net = d2l.EncoderDecoder(encoder, decoder)
  26. d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)

After the model is trained, we use it to translate a few English sentences into French and compute their BLEU scores.

```{.python .input}

@tab all

engs = [‘go .’, “i lost .”, ‘he\’s calm .’, ‘i\’m home .’] fras = [‘va !’, ‘j\’ai perdu .’, ‘il est calme .’, ‘je suis chez moi .’] for eng, fra in zip(engs, fras): translation, dec_attention_weight_seq = d2l.predict_seq2seq( net, eng, src_vocab, tgt_vocab, num_steps, device, True) print(f’{eng} => {translation}, ‘, f’bleu {d2l.bleu(translation, fra, k=2):.3f}’)

  2. ```{.python .input}
  3. #@tab all
  4. attention_weights = d2l.reshape(
  5.     d2l.concat([step[0][0][0] for step in dec_attention_weight_seq], 0),
  6.     (1, 1, -1, num_steps))

By visualizing the attention weights when translating the last English sentence, we can see that each query assigns non-uniform weights over key-value pairs. It shows that at each decoding step, different parts of the input sequences are selectively aggregated in the attention pooling.

```{.python .input}

Plus one to include the end-of-sequence token

d2l.show_heatmaps( attention_weights[:, :, :, :len(engs[-1].split()) + 1], xlabel=’Key posistions’, ylabel=’Query posistions’)

  2. ```{.python .input}
  3. #@tab pytorch
  4. # Plus one to include the end-of-sequence token
  5. d2l.show_heatmaps(
  6.     attention_weights[:, :, :, :len(engs[-1].split()) + 1].cpu(),
  7.     xlabel='Key posistions', ylabel='Query posistions')

Summary

  • When predicting a token, if not all the input tokens are relevant, the RNN encoder-decoder with Bahdanau attention selectively aggregates different parts of the input sequence. This is achieved by treating the context variable as an output of additive attention pooling.
  • In the RNN encoder-decoder, Bahdanau attention treats the decoder hidden state at the previous time step as the query, and the encoder hidden states at all the time steps as both the keys and values.

Exercises

  1. Replace GRU with LSTM in the experiment.
  2. Modify the experiment to replace the additive attention scoring function with the scaled dot-product. How does it influence the training efficiency?

:begin_tab:mxnet Discussions :end_tab:

:begin_tab:pytorch Discussions :end_tab: