All heads were trained on two NVIDIA 1080 Ti GPUs using the Pytorch deep learning library [1]}. For our backbone, we used pre-trained weights from the Pytorch Image Models library [2]}. For all experiments, we trained each head for 70 epochs, and used a regularization parameter of \(\beta = 2 \times 10^{-2}\) , the AdamW optimizer [3]}, and a cosine learning rate schedule with a warmup of three epochs with warmup learning rate of \(2 \times 10^{-7}\) , an initial learning rate chosen based on training accuracy of either \(5 \times 10^{-3}\)  or \(10^{-4}\) , and a final learning rate of \(2 \times 10^{-2}\)  times the initial learning rate. Data augmentation was performed using AutoAugment [4]}, along with color jittering, label smoothing, and training data interpolation. All heads aside from the self-attention head were trained using a batch size of 100, whereas the self-attention head was trained with a batch size of 20. For all ReLU heads, we choose a number of neurons such that the number of parameters across FNO, MLP-Mixer, self-attention, and MLPs are roughly equal. We provide information about the number of parameters in each head in Table REF .
