In principle, yes, you can, by using something like this on the last layer https://github.com/glassroom/heinsen_routing#sequence-to-vector https://github.com/gl

In principle, yes, you can, by using something like this on the last layer https://github.com/glassroom/heinsen_routing#sequence-to-vector https://github.com/glassroom/heinsen_routing#sequence-to-vec... ... but keep in mind it will not match the training efficiency of triangular masking in autoregressive transformers. You'd have to train the model one time-step at a time, vs. all time-steps in parallel with a masked autoregressive transformer. If you're looking to make context longer, a more practical approach is to route multiple past sequences of frozen pretrained transformer hidden states to a much shorter sequence, say, of length C, concatenate that with the current final sequence of hidden states, say, of length N, apply masked autoregressive attention on the C+N hidden states and then a feedforward layer to predict all next tokens in parallel, and last, drop the first C predictions, i.e., use only the trailing N predicted next tokens for training and inference.