Class token implementation

[karttikeya]

Hi,

I was wondering how the class token is supposed to be handled in the reversible design? Since, replicating the token across the two residual paths is perhaps not optimal.

Any thoughts/pointers to code is appreciated.

[lucidrains]

@karttikeya for the Reformer, I'd follow what this paper has done https://arxiv.org/abs/2103.17239 and only have the CLS token cross attend to the full sequence for about ~2-3 rounds at the end, as means of attention pooling