Skip to content

Extend featurewise transformations to layers in encoder

Emil Haines requested to merge ehaines/salt:ehaines-featurewise-encoder into main

Add ability to apply featurewise transformations (for conditioning the model output on some parameters) to the features of each layer in the encoder: see #57. This extends the current implementation in which transformations can only be applied to model inputs and the encoder output embeddings.

Also adds option to apply a layer normalisation to the transformed features (by setting apply_norm: True in the relevant featurewise_net configuration), as this could improve model stability and performance.

Review checklist:

  • CI Passing
  • Comments addressed
  • Source branch is up to date with target
Edited by Emil Haines

Merge request reports