  SGD performs worse than Adam by a significant margin on Transformers, but the
reason remains unclear. In this work, we provide an explanation through the
lens of Hessian: (i) Transformers are "heterogeneous": the Hessian spectrum
across parameter blocks vary dramatically, a phenomenon we call "block
heterogeneity"; (ii) Heterogeneity hampers SGD: SGD performs worse than Adam on
problems with block heterogeneity. To validate (i) and (ii), we check various
Transformers, CNNs, MLPs, and quadratic problems, and find that SGD can perform
on par with Adam on problems without block heterogeneity, but performs worse
than Adam when the heterogeneity exists. Our initial theoretical analysis
indicates that SGD performs worse because it applies one single learning rate
to all blocks, which cannot handle the heterogeneity among blocks. This
limitation could be ameliorated if we use coordinate-wise learning rates, as
designed in Adam.
