ReformerLM document representation about trax HOT 4 CLOSED

I think the hack is perfectly fine! Since it sounds like it worked, I'm closing the issue for now, but please re-open if more is needed.

from trax.

The state tree has usually the same form as the model tree, so state[-1] could be the state of the last layer. But the state does not contain activations -- usually (for most layers) it's just empty.

To have the last activations in state, you could just put them there in the forward_with_state method of the layer. We just today implemented a caching layer (for decoding) which does that, but only runs the layer once. What you want is something very similar, but simply running and putting the result in the state each time. Take a look at this code and let us know - I'm happy to help more!
https://github.com/google/trax/blob/master/trax/layers/combinators.py#L645

from trax.

Hi Lukasz,
Thank you for your reply.
I basically understood what you meant. Although I am experienced with TF, I am new to trax and was not able to pinpoint the location where I should be saving states.
Nevertheless, I thought of a quick workaround.
I created a new method called ReformerLMDocRep in which I deleted the dense and logsoftmax from the Map.
Naturally I had to delete the corresponding weights from the optimizer state that is loaded from a pretrained model.
This allowed me to return the output of the last layer of the model as "logits".
I imagine that this is the representation of the last token (and hence the document vector that I think I am looking for.)
Kindly let me know if you think that my hack is appropriate or not.
Of course if you give me the necessary instructions then I will be able to code up a better solution.
Thanks again for your help as I was able to better understand trax.

from trax.

Thanks for the confirmation. Ok to close.

from trax.