NMT Perceiver Encoder

[michalivne]

This PR separates the learning framework (i.e., loss) from the architecture, and adds a new Perceiver encoder bottleneck architecture.

Summary

A few major changes are introduced here:

1. Separating the learning framework (i.e., loss, or model type) from the architecture of the encoder / decoder.
2. The bottleneck encoder transformer can now be constructed using get_nemo_transformer method.
3. Supported learning frameworks:
   • NLL - Conditional cross entropy (the usual NMT loss)
   • VAE - Variational Auto-Encoder (https://arxiv.org/pdf/1312.6114.pdf)
   • MIM - Mutual Information Machine (https://arxiv.org/pdf/2003.02645.pdf)
4. Supported encoder architectures:
   • seq2seq - the usual transformer encoder without a bottleneck
   • bridge - attention bridge bottleneck (https://arxiv.org/pdf/1703.03130.pdf)
   • perceiver (new) - Perceiver bottleneck (https://arxiv.org/pdf/2103.03206.pdf)

The changes allow to easily select the architecture via configuration. Specifically, using a shared base YAML configuration and overriding flags during experimentation allows easy comparison of bottleneck and non-bottleneck architectures. Separating the loss from the architecture leads to a cleaner code which is easier to understand.

YAML Configuration

The following configurations in YAML config allows to control the learning framework and architecture:

model:
  ...
  model_type: 'nll' # learning (i.e., loss) type: nll (i.e., cross-entropy/auto-encoder), mim, vae (see description above)
  min_logv: -6 # minimal allowed log variance for mim
  latent_size: -1 # dimension of latent (projected from hidden) -1 will take value of hidden size
  non_recon_warmup_batches: 200000 # warm-up steps for mim, and vae losses
  recon_per_token: true # when false reconstruction is computed per sample, not per token
...
  encoder:
    library: nemo
    ...
    arch: seq2seq # seq2seq, bridge, perceiver (see description above)
    hidden_steps: 32 # fixed number of hidden steps
    hidden_blocks: 1 # number of repeat blocks (see classes for description)
    hidden_init_method: default # see classes for available values

  decoder:
    library: nemo
    ...
    arch: seq2seq # currently only seq2seq is supported

Detailed description of config parameters:

• model.encoder.arch=seq2seq
  • model.encoder.hidden_steps is ignored
  • model.encoder.hidden_blocks is ignored
  • model.encoder.hidden_init_method is ignored
• model.encoder.arch=bridge
  • model.encoder.hidden_steps: input is projected to the specified fixed steps
  • model.encoder.hidden_blocks: number of encoder blocks to repeat after attention bridge projection
  • model.encoder.hidden_init_method:
    • enc_shared (default) - apply encoder to inputs, than attention bridge, followed by hidden_blocks number of the same encoder (pre and post encoders share parameters)
    • identity - apply attention bridge to inputs, followed by hidden_blocks number of the same encoder
    • enc - similar to enc_shared but the initial encoder has independent parameters
• model.encoder.arch=perceiver
  • model.encoder.hidden_steps: input is projected to the specified fixed steps
  • model.encoder.hidden_blocks: number of cross-attention + self-attention blocks to repeat after initialization block (all self-attention and cross-attention share parameters)
  • model.encoder.hidden_init_method:
    • params (default) - hidden state is initialized with learned parameters followed by cross-attention with independent parameters
    • bridge - hidden state is initialized with an attention bridge

Usage Example

NOTE: enc_dec_nmt-bottleneck.py must be used. enc_dec_nmt.py does not support the required configuration to run properly and will likely fail.

python -- enc_dec_nmt-bottleneck.py \
      --config-path=conf \
      --config-name=aayn_bottleneck \
      ...
      model.model_type=nll \
      model.non_recon_warmup_batches=7500 \
      model.encoder.arch=perceiver \
      model.encoder.hidden_steps=32 \
      model.encoder.hidden_blocks=2 \
      model.encoder.hidden_init_method=params \
      ...

Additional Info

• Classes TransformerBottleneckEncoderNM, TransformerBottleneckDecoderNM, and corresponding config classes NeMoTransformerBottleneckConfig, NeMoTransformerBottleneckEncoderConfig, NeMoTransformerBottleneckDecoderConfig were added to nemo/collections/nlp/modules/common/transformer/transformer_bottleneck.py
• Class PerceiverEncoder was added to nemo/collections/nlp/modules/common/transformer/perceiver_encoders.py
• Class BridgeEncoder was added to nemo/collections/nlp/modules/common/transformer/bridge_encoders.py
• Classes TransformerBottleneckEncoderNM and TransformerEncoderNM test for src/tgt ids in the tokenizer valid range.

[lgtm-com]

This pull request introduces 1 alert when merging 53cce66 into 4abe5d5 - view on LGTM.com

new alerts:

• 1 for Unreachable code

[ericharper]

LGTM. Great PR! Thanks for all of the changes.

[MaximumEntropy]

LGTM. Thanks