NMT bottleneck

[michalivne]

This PR adds a bottleneck architecture to NMT models, along with support for training VAE and MIM latent variable models.

Summary

The bottleneck class (MTBottleneckModel) supports two architectures:

• No bottleneck: model_type=seq2seq, the usual NMT mdoel.
• Fixed-size bottleneck: model_type in [seq2seq-br, seq2seq-mim, seq2seq-vae], where the output of the encoder is projected to a fixed number of steps.

The projection to a fixed number of steps is based on the paper https://arxiv.org/pdf/1703.03130.pdf.
The idea is to use K attention heads to compute K weighted-average hidden states, projecting a variable number of steps into K steps.

The bottleneck variants offer different losses:

• The loss of seq2seq-br is reconstruction only (like seq2seq).
• The loss of seq2seq-mim is reconstruction + latent entropy minimization (See MIM, https://arxiv.org/pdf/2003.02645.pdf).
• The loss of seq2seq-vae is reconstruction + latent entropy regularization (See https://arxiv.org/pdf/1312.6114.pdf) correspondingly.

YAML Configuration

The following configurations were added to the YAML config:

model:
  model_type: 'seq2seq-br' # supports seq2seq, seq2seq-br, seq2seq-mim, seq2seq-vae (see description above)
  min_logv: -8 # minimal allowed logv for seq2seq-mim
  ortho_loss_coef: 0.0 # orthogonality coefficient for attention bridge
  att_bridge_size: 512 # dimension of a step in attention bridge
  att_bridge_k: 16 # fixed number of steps in attention bridge
  att_bridge_inner_size: 1024 # feedforward size in attention bridge
  non_recon_warmup_batches: 200000 # warm-up steps for seq2seq-mim, seq2seq-vae
  recon_per_token: true # when false reconstruction is computed per sample, not per token

• ortho_loss_coef - encou
• non_recon_warmup_batches - anneals the KL divergence term (for VAE) or latent entropy term (for MIM)
• recon_per_token - if false, loss is computed per sample (i.e., summed over all tokens), otherwise loss is averaged per token (the default behaviour).

Usage

See usage example below, training a seq2seq-br with 32 steps.

NOTE: max_generation_delta must be big enough to allow the generation of the longest sequence given the chosen bottleneck size.

python -- examples/nlp/machine_translation/enc_dec_nmt-bottleneck.py \
      --config-name=aayn_bottleneck \
      ...
      model.max_generation_delta=256 \
      ...
      model.model_type=seq2seq-br \
      model.att_bridge_size=1024 \
      model.ortho_loss_coef=0.0 \
      model.att_bridge_k=32 \
      model.att_bridge_inner_size=1024 \
      model.recon_per_token=true \
      model.non_recon_warmup_batches=150000 \
      ...

Additional Info

• The attention bridge class (AttentionBridge) was added to nemo/collections/nlp/modules/common/transformer/transformer_modules.py.
• The bottleneck class (MTBottleneckModel) supports logging of various loss terms.

[lgtm-com]

This pull request introduces 12 alerts when merging d3c655b into 44a3d02 - view on LGTM.com

new alerts:

• 12 for Unused import

[lgtm-com]

This pull request introduces 12 alerts when merging ec1303e into 44a3d02 - view on LGTM.com

new alerts:

• 12 for Unused import

[lgtm-com]

This pull request introduces 12 alerts when merging 78382bc into 44a3d02 - view on LGTM.com

new alerts:

• 12 for Unused import

[ericharper]

@michalivne, could you add usage instructions to the README?

[ericharper]

Please remove the empty file get_wk2.sh

[lgtm-com]

This pull request introduces 12 alerts when merging b91349b into 5363b49 - view on LGTM.com

new alerts:

• 12 for Unused import

[ericharper]

LGTM. Thanks!

[lgtm-com]

This pull request introduces 12 alerts when merging 14395e3 into 5363b49 - view on LGTM.com

new alerts:

• 12 for Unused import

[MaximumEntropy]

Thanks for making all of the changes :)

[lgtm-com]

This pull request introduces 12 alerts when merging 85449f3 into 5363b49 - view on LGTM.com

new alerts:

• 12 for Unused import

[lgtm-com]

This pull request introduces 12 alerts when merging efa5d3f into 6ebbcb8 - view on LGTM.com

new alerts:

• 12 for Unused import

[lgtm-com]

This pull request introduces 12 alerts when merging d95f97f into ed08545 - view on LGTM.com

new alerts:

• 12 for Unused import