Need Help with Improving Precision in Discourse Boundary Detection Model

qrm_dev: <p>Hello Pytorch Community,</p>
<p>I am currently working on a discourse segmentation project and have hit a roadblock. I’ve developed a model to classify tokens as either being a discourse boundary or not. The model is built using PyTorch Lightning and BERT.</p>
<p><strong>Issue:</strong> The main issue I’m facing is that the model consistently predicts 1s (boundaries), resulting in extremely low precision (&lt;0.001) and the recall is 1.0. This indicates a high number of false positives.</p>
<p><strong>Model Details:</strong> The model architecture is as follows:</p>
<ul>
<li>BERT for token embeddings.</li>
<li>A dropout layer (0.5).</li>
<li>A linear layer for classification.</li>
</ul>
<p>I am using the <code>F.logsigmoid</code> function for calculating log probabilities and <code>F.nll_loss</code> for loss calculation. Additionally, I’ve added L2 regularization.</p>
<pre><code class="lang-auto">class BoundaryClassification(pl.LightningModule):
    def __init__(self, bert):
        super().__init__()
        self.bert = bert
        self.dropout = nn.Dropout(0.5)
        self.segmentation_head = nn.Linear(self.bert.config.hidden_size, 1)

        
    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
        outputs = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
        sequence_output = outputs.last_hidden_state
        sequence_output = self.dropout(sequence_output)
        logits = self.segmentation_head(sequence_output)

        return logits

    def process_outputs(self, outputs):
        predictions = torch.argmax(outputs, dim=-1)
        y_pred = predictions.flatten().tolist()
   
        return y_pred

    def training_step(self, batch):
        input_ids = batch['input_ids']
        attention_mask = batch['attention_mask']
        labels = batch['edus']
        logits = self(input_ids=input_ids, attention_mask=attention_mask)

        log_probs = F.logsigmoid(logits)

        loss = F.nll_loss(log_probs.view(-1), labels.view(-1), weight=self.pos_weight)

        lambda_reg = 1e-4
        l2_reg = sum(p.pow(2.0).sum() for p in self.parameters())
        l2_reg = lambda_reg * l2_reg / 2


        probs = torch.sigmoid(logits).squeeze(2)

        y_pred = (probs &gt; 0.5).long().flatten().tolist()
        y_true = labels.cpu().flatten().detach().numpy().tolist()
        precision, recall = calculate_precision_recall(y_true, y_pred)
        pk = calculate_pk_score(y_true, y_pred)
        wd = calculate_windowdiff(y_true, y_pred)

        batch_size = len(labels)
        self.log("train_precision", precision, on_step=True, on_epoch=True, prog_bar=True, batch_size=batch_size)
        self.log('train_recall', recall, on_step=True, on_epoch=True, prog_bar=True, batch_size=batch_size)
        self.log('train_pk_score', pk, on_step=True, on_epoch=True, prog_bar=True, batch_size=batch_size)
        self.log('train_windowdiff', wd, on_step=True, on_epoch=True, prog_bar=True, batch_size=batch_size)
        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, batch_size=batch_size)

        return loss
</code></pre>
<p><strong>Specific Questions:</strong></p>
<ol>
<li>Could there be an issue with how I’m processing the logits or setting up the loss function?</li>
<li>Are there any specific techniques or adjustments in the model architecture that could help balance precision and recall?</li>
<li>Any insights on how to effectively handle class imbalance in this context?</li>
</ol>
<p>I appreciate any guidance or suggestions you can offer. Thank you for your time and help!</p>