reproducing https://arxiv.org/pdf/2201.02177 also tested https://github.com/ironjr/grokfast
TODOs
- clean up train script
- increase batch size
- try more fns. make new ones, and impl S5 and mod div fns
- larger train sets converging weirdly?
- grokfast loop is pretty inefficient, surely there's a better way?
- grokfast not working, either set up wrong or need to configure better
- is torch.nn.utils.clip_grad_norm_ standard? try removing