I have a doubt on I should set up the following problem:
Data:
My data is a tensor with shape (N, J, F) where N is the batch size, J is the sequence length, and F is the number of features of each element.
I will call each sequence of shapes (J, F) an "event" (the context is particle physics). (In reality, the length of the sequence is variable but I decided to fix the length and use the zero padding for the missing elements of the sequence)
(The order of the sequence is not important. No positional encoding is needed)
Task
For each event, I have to choose one element of the sequence, so the labels in an event are all zeros except for one 1.
This is just the first simple step, then I will have to do the same thing but for K classes (assign one element to the class '1', ..., assing one element to the class K and assign the rest of the sequence to the class 0).
But for now, just consider the 0/1 case.
What I have just done is build a self-attention network with a cross-entropy loss that has as output a softmax layer with shape (N, J, 1) and take as '1' the element of the sequence with the maximum score.
I think that this approach is essentially wrong because what I am asking to the model is to perform a multiclass classification with J classes, but this is not the problem that I want to solve.
Another idea is to build a softmax output with shape (N, J, 2) such that each element has a probability of being '0' or '1' and take the element with the highest probability of being '1'.
I don't know if this could be a good idea or if there is another method to address this problem. In the latter case, I am just binary classifying each element of the sequence, but this is not my task. I have to select one element inside a sequence
I think that the key is to get rid of the Cross-Entropy Loss and choose another Loss that adapts better to the problem.
Any advice is appreciated.
