@KFrank, thank you so much for taking time to write this answer. Really appreciate it. I would like to answer pointwise to your questions.
My xs are known before hand. It can be, for example, binary representation of 0 to 10 (each representation having length 4).
The reason for 10 classes is that the loss function that I use, ||y-ax||^2, is only accurate when x is the representation of one of the 10 classes.



 KFrank:

Softmax and MSELoss don’t really go together. Softmax would be
used for a classification (categorical) problem with a cross-entropy
loss (although you would actually leave out the Softmax and use
pytorch’s CrossEntropyLoss that has, in effect, Softmax built in).


This is exactly why I am on the hunt for another loss function that can help me train the network  .



 KFrank:

It looks to me as if your “vectors of length 4” make up batches of
batch size nBatch = 4


The sizes of y, a, x and b are all (batch_size, 4).
To summarize, from an optimization perspective, it can be shown that the loss function that I have mentioned in my question is optimal under this setting. Since the network has access to both a and y, I assume, the knowledge of both must be utilized for training the network. In essence, since the network has access to y and a and the all the possible classes (x), I believe the decision must based on the values of ||y-ax||^2. Thank you again!