Training a $2$-token attention is NP-hard

As the title indicates, this work is inspired by the celebrated paper of Blum and Rivest, "Training a $3$-node neural network is NP-complete." They showed that for a $2$-layer, $3$-node, $n$-input neural network (arguably the simplest neural network), deciding whether weights exist that fit a given training set exactly is NP-complete. We show that for an attention head with $1$-dimensional query, key, and value, and $n$-dimensional token input (arguably the simplest attention head), the analogous exact-fit problem on a training set of $2$-token sequences is NP-hard, and that if softmax is replaced with hardmax, the decision problem becomes NP-complete. These hardness results persist when fitting allows a fixed positive tolerance, measured coordinatewise or by mean squared loss.