# Multi-label token classification

**URL:** <https://discuss.huggingface.co/t/multi-label-token-classification/16509>\
**Category:** 🤗Transformers\
**Created:** [April 5, 2022, 7:12pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509 "2022-04-05T19:12:22Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![BunnyNoBugs](https://avatars.discourse-cdn.com/v4/letter/b/71c47a/32.png) [@BunnyNoBugs](https://discuss.huggingface.co/u/BunnyNoBugs)\
**Post date:** [April 5, 2022, 7:12pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/1 "2022-04-05T19:12:23Z")

</div>

Hi!

I am trying to solve a token classification problem in a multi-label setup. So far I haven’t found the best path to do it.  
Unlike `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` lacks the `problem_type` parameter, where you can specify that you’re working with multi-label.  
Should I deepen into logits and create a custom loss or is there a more straightforward solution?

Thanks!

---

<div class="post-metadata">

**Author:** ![drussellmrichie](https://avatars.discourse-cdn.com/v4/letter/d/e5b9ba/32.png) [@drussellmrichie](https://discuss.huggingface.co/u/drussellmrichie)\
**Post date:** [April 11, 2022, 9:19pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/2 "2022-04-11T21:19:52Z")

</div>

I would also love to be able to do this. I’m trying to do some NER that has multiple layers of annotation such that a given token could have more than one label…

---

<div class="post-metadata">

**Author:** ![BunnyNoBugs](https://avatars.discourse-cdn.com/v4/letter/b/71c47a/32.png) [@BunnyNoBugs](https://discuss.huggingface.co/u/BunnyNoBugs)\
**Post date:** [April 11, 2022, 10:05pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/3 "2022-04-11T22:05:47Z")

</div>

Actually, I already started implementing a custom loss within the `compute_loss` function of the `Trainer`…

…but the thing surely deserves a feature request on GitHub! Will try to do that

---

<div class="post-metadata">

**Author:** ![drussellmrichie](https://avatars.discourse-cdn.com/v4/letter/d/e5b9ba/32.png) [@drussellmrichie](https://discuss.huggingface.co/u/drussellmrichie)\
**Post date:** [April 12, 2022, 4:15am UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/4 "2022-04-12T04:15:09Z")

</div>

Oh nice. Do you by any chance have code you could share? 😉

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [April 12, 2022, 8:01am UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/5 "2022-04-12T08:01:11Z")

</div>

I don’t think it should be much harder than swapping out `CrossEntropyLoss` for `BCEWithLogitsLoss` and one-hot encoding your labels and making sure that they are floats.

---

<div class="post-metadata">

**Author:** ![BunnyNoBugs](https://avatars.discourse-cdn.com/v4/letter/b/71c47a/32.png) [@BunnyNoBugs](https://discuss.huggingface.co/u/BunnyNoBugs)\
**Post date:** [April 12, 2022, 8:28pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/6 "2022-04-12T20:28:44Z")

</div>

Yeah, that’s roughly what I am doing. Will share the code once it’s done!

---

<div class="post-metadata">

**Author:** ![drussellmrichie](https://avatars.discourse-cdn.com/v4/letter/d/e5b9ba/32.png) [@drussellmrichie](https://discuss.huggingface.co/u/drussellmrichie)\
**Post date:** [April 19, 2022, 1:26pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/7 "2022-04-19T13:26:02Z")

</div>

@BunnyNoBugs Any update on this? 🥺 I have been working on this myself but am running into some challenges…

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [April 19, 2022, 1:38pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/8 "2022-04-19T13:38:33Z")

</div>

This is an untested attempt, but I think it should work. Read more about BCEWithLogitsLoss [here](https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html).

```python
class MultiLabelTrainer(Trainer):
    def __init__ (self, *args, class_weights: Optional[FloatTensor] = None, **kwargs):
        super(). __init__ (*args, **kwargs)
        if class_weights is not None:
            class_weights = class_weights.to(self.args.device)
            logging.info(f"Using multi-label classification with class weights", class_weights)
        self.loss_fct = BCEWithLogitsLoss(weight=class_weights)

    def compute_loss(self, model, inputs, return_outputs=False):
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.
        Subclass and override for custom behavior.
        """
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        try:
            loss = self.loss_fct(outputs.logits.view(-1, model.num_labels), labels.view(-1))
        except AttributeError: # DataParallel
            loss = self.loss_fct(outputs.logits.view(-1, model.module.num_labels), labels.view(-1))

        return (loss, outputs) if return_outputs else loss

```

Make sure your data is correctly formatted, e.g., 0 or 1 encoded for each label for each sample.

---

<div class="post-metadata">

**Author:** ![BunnyNoBugs](https://avatars.discourse-cdn.com/v4/letter/b/71c47a/32.png) [@BunnyNoBugs](https://discuss.huggingface.co/u/BunnyNoBugs)\
**Post date:** [April 19, 2022, 1:58pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/9 "2022-04-19T13:58:49Z")

</div>

Yep, I did the same but without `view` and getting an error with the dims. Will get back to this task on Thursday

---

<div class="post-metadata">

**Author:** ![drussellmrichie](https://avatars.discourse-cdn.com/v4/letter/d/e5b9ba/32.png) [@drussellmrichie](https://discuss.huggingface.co/u/drussellmrichie)\
**Post date:** [April 19, 2022, 3:52pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/10 "2022-04-19T15:52:48Z")

</div>

Thanks so much @BunnyNoBugs and @BramVanroy! I’m trying this now and getting some other issues but I think they are specific to my situation. I may report back later with an update if I think it will be useful to others. Thanks again!

---

<div class="post-metadata">

**Author:** ![drussellmrichie](https://avatars.discourse-cdn.com/v4/letter/d/e5b9ba/32.png) [@drussellmrichie](https://discuss.huggingface.co/u/drussellmrichie)\
**Post date:** [April 20, 2022, 7:46pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/11 "2022-04-20T19:46:20Z")

</div>

Here’s a question: doesn’t the custom loss function need to ignore the predictions for special tokens like CLS, PAD, SEP, and (if one is only applying labels to the first subword in a token) the non-first subwords?

---

<div class="post-metadata">

**Author:** ![aligator](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/aligator/32/4534_2.png) [@aligator](https://discuss.huggingface.co/u/aligator)\
**Post date:** [April 21, 2022, 9:10am UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/12 "2022-04-21T09:10:00Z")

</div>

Thanks this was super helpful! I followed your tips and got dimension errors, but fixed it by modifying the view parameters, and I also had to cast the labels to floats:

```auto
loss = self.loss_fct(outputs.logits.view(-1, model.num_labels), labels.view(-1, model.num_labels).float())

```

---

<div class="post-metadata">

**Author:** ![drussellmrichie](https://avatars.discourse-cdn.com/v4/letter/d/e5b9ba/32.png) [@drussellmrichie](https://discuss.huggingface.co/u/drussellmrichie)\
**Post date:** [April 21, 2022, 1:59pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/13 "2022-04-21T13:59:26Z")

</div>

I need to look at my results more closely but I think I’ve got this working. My model only gets to about 60% F1 but on the validation set but I think I have ideas about how to improve it. 👿

```auto
class MultiLabelNERTrainer(Trainer):
    def __init__ (self, *args, class_weights: Optional[FloatTensor] = None, **kwargs):
        super(). __init__ (*args, **kwargs)
        if class_weights is not None:
            class_weights = class_weights.to(self.args.device)
            logging.info(f"Using multi-label classification with class weights", class_weights)
        self.loss_fct = BCEWithLogitsLoss(weight=class_weights)

    def compute_loss(self, model, inputs, return_outputs=False):
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.
        Subclass and override for custom behavior.
        """
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        
        # this accesses predictions for tokens that aren't CLS, PAD, or the 2nd+ subword in a word
        # and simultaneously flattens the logits or labels
        flat_outputs = outputs.logits[labels!=-100] 
        flat_labels = labels[labels!=-100]
        
        try:
            loss = self.loss_fct(flat_outputs, flat_labels)
        except AttributeError: # DataParallel
            loss = self.loss_fct(flat_outputs, flat_labels)

        return (loss, outputs) if return_outputs else loss

```

---

<div class="post-metadata">

**Author:** ![BunnyNoBugs](https://avatars.discourse-cdn.com/v4/letter/b/71c47a/32.png) [@BunnyNoBugs](https://discuss.huggingface.co/u/BunnyNoBugs)\
**Post date:** [April 21, 2022, 8:46pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/14 "2022-04-21T20:46:22Z")

</div>

It is said, for example, [here](https://huggingface.co/transformers/v4.9.2/main_classes/data_collator.html#transformers.data.data_collator.DataCollatorForTokenClassification), that -100 value is automatically ignored by PyTorch loss functions.

---

<div class="post-metadata">

**Author:** ![BramVanroy](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/bramvanroy/32/47360_2.png) [@BramVanroy](https://discuss.huggingface.co/u/BramVanroy)\
**Post date:** [April 22, 2022, 8:41am UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/15 "2022-04-22T08:41:43Z")

</div>

That’s only true for CrossEntropyLoss AFAIK. You can compare the signatures. [CrossEntropyLoss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) has an “ignore\_index” option, [BCEWithLogitsLoss](https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html) does not. So I think that @drussellmrichie 's adaptation is indeed needed.

---

<div class="post-metadata">

**Author:** ![BunnyNoBugs](https://avatars.discourse-cdn.com/v4/letter/b/71c47a/32.png) [@BunnyNoBugs](https://discuss.huggingface.co/u/BunnyNoBugs)\
**Post date:** [April 22, 2022, 2:17pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/16 "2022-04-22T14:17:13Z")

</div>

I haven’t finished my loss function yet, but here’s what I found out:

1. `inputs.pop('labels')` in @BramVanroy’s example is very important: otherwise the labels are passed to the model and the standard `CrossEntropyLoss` inside it is computed instead of the custom one.

2. We should make use of the `attention_mask` in the custom loss. However, it may be the same as @drussellmrichie’s suggestion about the special tokens, maybe someone will correct me.

---

<div class="post-metadata">

**Author:** ![drussellmrichie](https://avatars.discourse-cdn.com/v4/letter/d/e5b9ba/32.png) [@drussellmrichie](https://discuss.huggingface.co/u/drussellmrichie)\
**Post date:** [April 22, 2022, 2:45pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/17 "2022-04-22T14:45:07Z")

</div>

> 1. We should make use of the `attention_mask` in the custom loss. However, it may be the same as @drussellmrichie’s suggestion about the special tokens, maybe someone will correct me.

I’m fairly confident that it’s not necessary given the chunk I added with `labels!=-100`.

EDIT: Just to elaborate: the proof is in the pudding, and indeed I was able to train some decent multi-label NER models with this approach. 🙂

---

<div class="post-metadata">

**Author:** ![BunnyNoBugs](https://avatars.discourse-cdn.com/v4/letter/b/71c47a/32.png) [@BunnyNoBugs](https://discuss.huggingface.co/u/BunnyNoBugs)\
**Post date:** [April 22, 2022, 6:36pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/18 "2022-04-22T18:36:35Z")

</div>

What do you mean with padding?

---

<div class="post-metadata">

**Author:** ![drussellmrichie](https://avatars.discourse-cdn.com/v4/letter/d/e5b9ba/32.png) [@drussellmrichie](https://discuss.huggingface.co/u/drussellmrichie)\
**Post date:** [April 22, 2022, 7:09pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/19 "2022-04-22T19:09:42Z")

</div>

Are you referring to my phrase “the proof is in the pudding”? If so, I apologize – I am using an old proverb. 😉 I am just saying that the proof that `labels!=100` is doing something similar to attention mask lies in the fact that the model performs decently and generates reasonable predictions for the tokens we care about (not CLS, PAD, etc). Also, I noticed that the loss went way up after I introduced `labels!=100`, I think because now the easy predictions of the CLS, PAD, and 2nd+ subword tokens can no longer drive the loss down – that is consistent with @BramVanroy 's comment that [BCEWithLogitsLoss](https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html) does not ignore predictions for special tokens, and needed something like my `labels!=100` trick to ignore those.

Hope that clarifies.

---

<div class="post-metadata">

**Author:** ![murdockthedude](https://avatars.discourse-cdn.com/v4/letter/m/51bf81/32.png) [@murdockthedude](https://discuss.huggingface.co/u/murdockthedude)\
**Post date:** [May 17, 2022, 5:31pm UTC](https://discuss.huggingface.co/t/multi-label-token-classification/16509/20 "2022-05-17T17:31:33Z")

</div>

Hi @drussellmrichie (and @BunnyNoBugs and @BramVanroy !), thanks for the sample code above. A couple questions if you don’t mind:

- Is the custom trainer all that’s needed to adapt the model to multi-label classification?
- Also, is this being implemented on top of `AutoModelForTokenClassification`?
- Would you be willing to share an example notebook? I’d love to see the full implementation, including the “right” way to extract the predictions at inference time.

If anyone would be willing to share an example notebook, I’m currently trying to get this working and would very much appreciate a sample implementation!

Thanks!

[Next page](https://discuss.huggingface.co/t/multi-label-token-classification/16509.md?page=2)