I'm working on a GNN project associated with molecule classification.
The project is to classify if the atom in the molecule will initiate a certain reaction.
For example, a molecule can be represented as
[0, 1, 2, 3, 4], where the numbers indicate the index of the atoms. If the reaction happens on index 0 and 3, the
y_label will be
[1, 0, 0, 1, 0]. In reality, the
x_feature will have the shape
y_label will have the shape
I'm using two approaches to solve this problem.
- using softmax as output layer with MSE as the loss function
- using binary cross-entropy with logits as the loss function
In case 1, I'm thinking to have the probability as an output on each atom, where the probability indicates the probability of reaction for each atom. Then, I'm picking the top 3 atoms that have the highest probability as my reaction atoms.
In case 2, I'm thinking to treat each atom as a binary classification task, and if the output > 0.5, then this indicates the reaction will happen on this atom.
After experimentation, case 1 is much better than case 2, but I'm facing difficulty in interpreting why it's so.
Especially in case 1, can I use the softmax combined with the MSE to solve such a problem? I was expecting the MSE + softmax configuration to have a worse result than the binary cross-entropy one, but now I'm confused.
Can anyone help me to explain why this would happen?