Thanks for the replies @KFrank @googlebot!



 KFrank:

In fact, including x**2 lets you train an almost trivial model that
consists of a single Linear layer (and no activation functions):


Indeed a  simple Linear module with one layer does a much better job at finding the polynomial function than my original model. The beauty is in simplicity apparently.



 KFrank:

So the network has to spend many training iterations learning to ignore x**3 (or to process it into some approximation to a quadratic), which wastes valuable
time, and potentially sends the training off on some wild goose chase
from which it never recovers.


As both of you mentioned this seems to be cause of the problem I ran into in my original post. I’m still a bit surprised the network doesn’t figure out it simply needs to ignore x**3 though.



 KFrank:

A minor comment: MSELoss (using its defaults) takes the mean
of the loss over the batch, so loss / data_size is incorrectly
dividing by the batch size a second time.


Oops! Thanks for pointing that out, some leftover code from a previous version.



 KFrank:

I’d like to expand on this a little. It is true that Linear layers introduce
no quadratic terms. Also the ReLU activation function is piecewise
linear. But many non-linear activation functions do have quadratic (and
higher-order) terms in their expansions, so they do introduce quadratic
terms into the overall function computed by the network.


So in general, if one suspects that the data contains a certain relation/function, would you add said function to the input instead of letting the Neural Network try to find it?