Loss function for comparing two vectors for categorization

Viewed 522

I am performing a NLP task where I analyze a document and classify it into one of six categories. However, I do this operation at three different time periods. So the final output is an array of three integers (sparse), where each integer is the category 0-5. So a label looks like this: [1, 4, 5].

I am using BERT and am trying to decide what type of head I should attach to it, as well as what type of loss function I should use. Would it make sense to use BERT's output of size 1024 and run it through a Dense layer with 18 neurons, then reshape into something of size (3,6)?

Finally, I assume I would use Sparse Categorical Cross-Entropy as my loss function?

2 Answers

The bert final hidden state is (512,1024). You can either take the first token which is the CLS token or take the average pooling. Either way your final output is shape (1024,) now simply put 3 linear layers of shape (1024,6) as in nn.Linear(1024,6) and pass it into the loss function below. (you can make it more complex if you want to)

Simply add up the loss and call backward. Remember you can call loss.backward() on any scalar tensor.(pytorch)

def loss(time1output,time2output,time3output,time1label,time2label,time3label):
    loss1 = nn.CrossEntropyLoss()(time1output,time1label)
    loss2 = nn.CrossEntropyLoss()(time2output,time2label)
    loss3 = nn.CrossEntropyLoss()(time3output,time3label)

    return loss1 + loss2 + loss3

In a typical setup you take a CLS output of BERT (a vector of length 768 in case of bert-base and 1024 in case of bert-large) and add a classification head (it may be a simple Dense layer with dropout). In this case the inputs are word tokens and the output of the classification head is a vector of logits for each class, and usually a regular Cross-Entropy loss function is used. Then you apply softmax to it and get probability-like scores for each class, or if you apply argmax you will get the winning class. So the result might be either vector of classification scores [1x6] or the dominant class index (an integer).

enter image description here Image taken from d2l.ai

You can simply concatenate 3 such networks (for each time period) to get the desired result.

Obviously, I have described only one possible solution. But as it is usually provide good results I suggest you try it before moving over to more complex ones.

Finally, Sparse Categorical Cross-Entropy loss is used when output is sparse (say [4]) and regular Categorical Cross-Entropy loss is used when output is one-hot encoded (say [0 0 0 0 1 0]). Otherwise they are absolutely the same.

Related