NMT Perceiver Encoder
State: closed | Pull request: yes

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.