SoftmaxLayer with one-hot vectors

The quantum physics problem I am looking at is a rotated qubit. The setup involves a y-rotated qubit measured in the z-basis (hence spin-up and spin-down). This scenario typically includes modeling the quantum state of the qubit and analytically determining the measurement outcome probabilities, then generating measurement outcomes for training for particular rotations and using another set of test measurement data to infer the most probable rotation angle. So I discretize the measurement angles and generate measurement results for each discrete angle to use as training data. I plan to generate another set of measurement results for testing later. My training data involves associating elements of the form $${1, 1, 1, 0, 1, 1, 0, 1, 0, 1} -> {0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}$$ where the measurements on the left are 1 if spin up and 0 if spin down, and the lists on the right are one-hot vector lists which correspond to the distinct discrete rotation angles from $$[0, \pi]$$.

Query: I am interested in using the SoftmaxLayer[], but I am getting an error when including it my NetTrain[] function. The error states: "Batch #1 will be skipped, because one or or more inputs provided to port "Output" was invalid: input is not an integer between 1 and 20. This batch will be ignored in subsequent training rounds. More information can be obtained via the "SkippedTrainingData" property."

The code runs if I remove the SoftmaxLayer[] from the NetTrain[] command. Can anyone advise on how to resolve this problem, many thanks for any assistance.

(*Spin-up/spin-down measurement probabilities on rotated qubit*)

Pu[\[Theta]_] := (Cos[\[Theta]/2])^2;
Pd[\[Theta]_] := 1 - Pu[\[Theta]];

(*Discretize rotation angles*)

theta = Table[\[Theta], {\[Theta], \[Pi]/20, \[Pi], \[Pi]/20}]

(*Form One-hot vectors from grid theta values*)
numBins = Length[theta];
binToOneHot[bins_List, totalBins_Integer] :=
Table[If[i == #, 1, 0], {i, 1, totalBins}] & /@ bins
bins = Range[numBins];
oneHotEncoded = binToOneHot[bins, Length[theta]]

(*Compute measurement probabilities for each grid theta value*)
pr = Table[
List[N[FullSimplify[Re[Pu[\[Theta]]]]],
N[FullSimplify[Re[Pd[\[Theta]]]]]], {\[Theta], \[Pi]/
20, \[Pi], \[Pi]/20}]

(*Generate measurement outcomes for each discrete angle (with random \
number generator)*)
m = 10;
meas = Table[{}, {Length[pr]}];
Table[Do[b = ConstantArray[0, 1]; Clear[r];
updateTally[] := Block[{r}, r = RandomReal[];
If[r < pr[[j]][[1]], b[[1]] += 1,
If[pr[[j]][[1]] <= r < pr[[j]][[1]] + pr[[j]][[2]],
b[[1]] += 0]]];
Do[updateTally[], {1}]; AppendTo[meas[[j]], b], {i, 1, m}], {j,
Length[pr]}];
meas = Partition[Flatten[meas], 10]
trainingData = Rule @@@ Transpose[{meas, oneHotEncoded}]

(*Neural network architecture*)
net = NetChain[{LinearLayer[4], ElementwiseLayer["ReLU"],
LinearLayer[20], SoftmaxLayer[]}];

(*Training the network*)
trainedNet = NetTrain[net, trainingData, MaxTrainingRounds -> 1]
$$$$

• is this helpful? Changing NetTrain's loss function to: trainedNet = NetTrain[net, trainingData, MaxTrainingRounds -> 1, LossFunction -> CrossEntropyLossLayer["Binary"]] allows the training to go without any messages. I'm not posting this as an answer because I don't know if this is the appropriate loss function to use here. But the answer to the question I linked above appears to explain why this issue occurs when SoftmaxLayer` is the last layer in a net
– ydd
Commented May 11 at 16:18
• @ydd It certainly is an appropriate loss function and your comment is very helpful (feel free to post it as a full answer). Many thanks. Commented May 12 at 0:49
• I think I will just leave it as a comment as I'm worried my answer will be too similar to the linked answer. I'm glad I could help though :D
– ydd
Commented May 13 at 21:57
• @ydd Okay understood. If you have a chance please check out this post. The accuracy of the neural network on the test data does not seems to be great, although the code works, maybe you can see how to improve it. Thanks for your time. Commented May 14 at 8:55