Understanding the moving parts for training

[nbgundavarapu]

Hi team!

The paper and the accompanying codebase are really great! We are trying to use S4 for a different problem and there are a lot of engineering details that seem to be affecting the training:

1. Adding/removing GLU activation => Also affects ListOps a lot
2. Dropout vs Dropout2D
3. Learning rate of the state space parameters and their schedule in relation to rest of parameters
4. Whether to share the diagonal and low-rank params depth-wise or not
5. Whether to share log_step depth-wise or not => Leads to NaN loss at times
6. Whether to use the NPLR formulation from S4 paper or Sashimi
7. Whether to use S4 or S4D or DSS
8. Bidirectional or unidirectional ...

From your experience, could you share any ideas on how to choose from these options? Also, could you list any other important details that might affect training?