Thank you @anantguptadbl!!
Just wanted to follow up here; I’ve only tried your second suggestion so far (changing batch size to 32), but it seems to lead to much better results. Really do appreciate your help; was stuck on this for a little while! Probably should’ve tried some sort of sweep over hyperparameters. Anyway, I’ve only ran this for a few hours so I’m sure it will get much better results later.
For anyone looking at this in the future:
Losses / Validation Losses

image3700×1210 337 KB

Histogram of values for intermediate values (memory = encoder output, predicted embedding = decoder output)

image3732×1194 215 KB

Last Attention Layers
Screen Shot 2022-02-02 at 8.34.41 AM

image1664×1618 221 KB

Screen Shot 2022-02-02 at 8.35.04 AM

image1758×1644 230 KB

Will follow up with a more detailed write up via a blog post s.t. others can replicate later. Appreciate your help again.