> it's brute forcing a small grid
If by small grid you are referring to the attention matrix plot shown, then that is not a correct interpretation. That diagonal-like pattern it learns, is 3x3 convolution, so it can compare the neighbours of a given cell.
Edit: and note that every grid it is trained on / runs inference on is randomly generated and completely unique, so it cannot just memorise examples