Why does the BERT NSP head linear layer have two outputs?

Question

Here's the code in question.

https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_bert.py#L491

class BertOnlyNSPHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, pooled_output):
        seq_relationship_score = self.seq_relationship(pooled_output)
        return seq_relationship_score

I think it was just ranking how likely one sentence would follow another? Wouldn't it be one score?

Answer 1

This seems to be inherited from the original Google implementation, which also uses 2 outputs (https://github.com/google-research/bert/blob/master/run_pretraining.py#L293). I can see two possible reasons that the original implementation uses 2 outputs:

1. They are using the cross entropy loss, which typically works with log probabilities. To get probabilities they use softmax activation, which requires an output for each class. It is possible, of course, to compute cross entropy from sigmoid activations (which would correspond to a 1-output architecture), but there seems to be some confusion as to whether the output of the sigmoid function should be used as a probability.

2. Using 2 outputs can simplify the computation of the binary cross entropy loss, which, in typical Google fashion, is computed using low-level tensorflow ops rather than with tf.nn.softmax_cross_entropy_with_logits. Specifically,

-tf.reduce_sum(one_hot_labels * log_probs, axis=-1)

where one_hot_labels and log_probs are $\mathbb{R}^{N \times 2}$, is much easier to read than

-tf.reduce_sum(binary_labels * tf.math.log(probs) + (1 - binary_labels) + tf.math.log(1 - probs))

where binary_labels and probs are $\mathbb{R}^N$.