Transformer Hydranets and Multi-Task Learning with Hugging Face and PyTorch: Ensembling vs. Entwinement

The Power of Multi-Task Learning

Multi-task learning mirrors human learning. Knowledge from one task aids in mastering others. Just as a child uses facial recognition skills to identify other objects, MTL applies learned insights across various tasks. In a more theoretical view MTL is a form of inductive transfer which entices models to favor hypotheses that explain multiple tasks, leading to better generalization.⁶ Here are a few perspectives summarized from Ruder (2017).

Use-Case for a Transformer Hydranet

Hydranets’ true power shines when there is overlapping knowledge that can be exploited between tasks. Lets consider an example of a content moderation solution. Traditional approaches might deploy separate models for different aspects of content moderation—such as sentiment analysis, toxic comment detection, and spam identification. However, a transformer hydranet can be trained to tackle all these tasks concurrently.

Sharing a common backbone means hydranet models can gain a more holistic understanding of language nuances crucial to discern both context and intent in content. The model can learn to simultaneously identify toxic language, categorize comments, and evaluate sentiment. Cues learned from each of these tasks can help the model succeed at other tasks in turn. In this sense, the overlapping knowledge can be exploited by the hydranet to create synergy during learning and thus more powerful backbones.

Hydranet models can also reduce compute resources for inference, especially in cases where they replace ensembles or iterate on top of single-task models. Most of the inference workload is spent computing the forward pass of the model backbone. The forward pass of each model head can be computed in parallel and its usually much lighter than the backbone.

Multi-label classification is another great use-case as it has been traditionally solved with ensembles like one-to-rest or target transformations like taking the power set all labels. In the multi-task learning paradigm, multi-label classification can be modeled directly with a hydranet approach. Rather than running multiple models or exploding the number of targets, we build the forward pass of the model to compute in parallel each label set—a much more organic approach. This works well when you have a strict limit on the overlap of classes, but can fail on pick-all-applicable scenarios.

Off The Shelf Transformer Heads

One of the great features of Hugging Face’s transformers library is the collection of pre-built heads for performing versatile tasks on top of pre-trained base models.⁷ When you navigate to a specific model in the library, like RoBERTa, you will see a set of various tasks that the model is capable of performing such as masked language modeling (MLM), token classification, and sequence classification.

Predefined task-specific models can be pulled down in a single line of code. Here’s how we can easily instantiate RoBERTa for both token and sequence classification tasks:

from transformers import RobertaForTokenClassification, RobertaForSequenceClassification

token_model = RobertaForTokenClassification.from_pretrained(
  'roberta-base', num_labels=5
)
sequence_model = RobertaForSequenceClassification.from_pretrained(
  'roberta-base', num_labels=3
)

In this snippet, num_labels is adjusted based on the task requirements, such as the number of entity types in Named Entity Recognition (NER) or classes in sequence (text) classification tasks like sentiment analysis.

Under the hood, transformers defines an explicit model architecture for each task, but in a spectacularly modular way.

def forward(self, input_ids, attention_mask=None, labels=None):
    # Run through the model backbone 
    outputs = self.roberta(input_ids, attention_mask=attention_mask)
    # Access the last hidden states (embeddings from final base model layer)
    sequence_output = outputs[0]
    # Feed the embeddings to the task-specific head
    logits = self.classifier(sequence_output)

    # Additional processing for loss calculation and output
    # ...
    
    return logits

class RobertaForTokenClassification(RobertaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        # Model backbone
        self.roberta = RobertaModel(config, add_pooling_layer=False)
        # Classification head (a simple MLP)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

class RobertaForSequenceClassification(RobertaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        # Model backbone
        self.roberta = RobertaModel(config, add_pooling_layer=False)
        # Classification head ([CLS] pooling + a simple MLP)
        self.classifier = RobertaClassificationHead(config)


class RobertaClassificationHead(nn.Module):
    """Head for sentence-level classification tasks."""

    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, features, **kwargs):
        # Take <s> token (equivalent to [CLS] token) 
        x = features[:, 0, :]
        x = self.dense(x)
        x = torch.tanh(x)
        x = self.out_proj(x)
        return x

class RobertaForQuestionAnswering(RobertaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        # Model backbone
        self.roberta = RobertaModel(config, add_pooling_layer=False)
        # QA head (a simple MLP identical to token classification head)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
        
    def forward(self, input_ids, attention_mask=None, labels=None):
        # Run through the model backbone 
        outputs = self.roberta(input_ids, attention_mask=attention_mask)
        # Access the last hidden states (embeddings from final base model layer)
        sequence_output = outputs[0]
        # Feed the embeddings to the task-specific head
        logits = self.qa_outputs(sequence_output) ## [batch_size, sequence_length, 2]
    
        ## QA Specific Differences ##
        # Split the logits tensor into two tensors along the last dimension.  
        start_logits, end_logits = logits.split(1, dim=-1) ## [batch_size, sequence_length, 1], [batch_size, sequence_length, 1]
        
        # Squeeze along the last dimension 
        start_logits = start_logits.squeeze(-1).contiguous() ## [batch_size, sequence_length]
        end_logits = end_logits.squeeze(-1).contiguous() ## [batch_size, sequence_length]
      
        return start_logits, end_logits

loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
start_loss = loss_fct(start_logits, start_positions)
end_loss = loss_fct(end_logits, end_positions)
total_loss = (start_loss + end_loss) / 2

Off With Their (Transformer) Heads!

Building a hydranet is one way to force yourself to recognize the beauty of PyTorch Autograd. The essential feature of hydranet architectures is that your model’s forward pass returns multiple objects.

Let’s take our base RoBERTa model and build a hydranet designed for content moderation. We want to be able to:

1. Flag harmful comments (harm, not harm)
2. Classify type of harm (e.g. spam, hate, promotions)

Both tasks can be modeled as sequence classification. Let’s start by stealing some code like all good engineers do! I’ll copy the source code for RobertaForSequenceClassification and we will build out the first head to perform binary sequence classification and the second head perform multi-class classification.¹² I’ve renamed a few things and added separate configs as the output dimensions are different between head 1 and 2:¹³

class ContentModerationHydraNet(RobertaPreTrainedModel):
    def __init__(self, config, harm_flag_config, harm_type_config):
        super().__init__(config)
        # Model backbone
        self.roberta = RobertaModel(config, add_pooling_layer=False)
        # Output a single logit for harm confidence
        self.harm_flag_classifier = RobertaClassificationHead(harm_flag_config)
        # Output 3 logits for spam, hate, promotion classes
        self.harm_type_classifier = RobertaClassificationHead(harm_type_config)
    
    def forward(self, input_ids, attention_mask=None, labels=None):
        # Run through the model backbone 
        outputs = self.roberta(input_ids, attention_mask=attention_mask)
        # Access the last hidden states (embeddings from final base model layer)
        sequence_output = outputs[0]
        # Feed the embeddings to head 1
        logits_harm_flag = self.harm_flag_classifier(sequence_output)
        # Feed the embeddings to head 2
        logits_harm_type = self.harm_type_classifier(sequence_output)

        # We now return a tuple of output for each head
        return logits_harm_flag, logits_harm_type

And that its, now we’ve built a transformer hydranet!

Advanced Hydranet Architecture

Let’s ratchet this up to the next level. To flag harmful comments we certainly need to understand the content of the comment, but we might also want to exploit additional context by using information about the user who makes the comment or the original post the comment is about. However, our current model can only take text as input.

We can inject additional information into our model in several interesting ways. If we have historical information on this user, it might be something useful for our harm_flag_classifier to learn on. We can leverage the user’s age and amount of comments that have been flagged for harm in the past as an example. To do this, simply concatenate these onto the final embeddings from the model backbone to provide additional signal:¹⁴

class ContentModerationHydraNet(RobertaPreTrainedModel):
    def __init__(self, config, harm_flag_config, harm_type_config):
        # Same as above
    
    # Add additional arguments for user-level features
    def forward(self, input_ids, age, prev_flags, attention_mask=None, labels=None):
        # Run through the model backbone 
        outputs = self.roberta(input_ids, attention_mask=attention_mask)
        # Access the last hidden states (embeddings from final base model layer)
        sequence_output = outputs[0]
        # Concatenate user-level features for head 1
        harm_flag_input = torch.cat((sequence_output, age.view(-1, 1), prev_flags.view(-1, 1)))
        # Feed the RoBERTa embeddings + user-level features to head 1
        logits_harm_flag = self.harm_flag_classifier(harm_flag_input)
        # Feed the RoBERTa embeddings to head 2
        logits_harm_type = self.harm_type_classifier(sequence_output)

        # We now return a tuple with output for each head
        return logits_harm_flag, logits_harm_type

Now, we are injecting additional features only into one of the model’s heads. That information can be used to better train the classifier but its also possible for it propagate all the way through to the backbone as well.

Let’s go one step further—we haven’t touched our harm_type_classifier yet. Since we are already classifying something as harmful or not with the harm_flag_classifier we might want to use that output to augment the harm_type_classifier as well. That means we will want to feed the output of the first head into the second head. It’s also possible to gate the execution of the second-head on the first head. For example, if we don’t flag a comment as harmful, we don’t need to pass it to the harm_type_classifier at all.

class ContentModerationHydraNet(RobertaPreTrainedModel):
    def __init__(self, config, harm_flag_config, harm_type_config):
        # Same as above
    
    # Add additional arguments for user-level features
    def forward(self, input_ids, age, prev_flags, attention_mask=None, labels=None):
        # Run through the model backbone 
        outputs = self.roberta(input_ids, attention_mask=attention_mask)
        # Access the last hidden states (embeddings from final base model layer)
        sequence_output = outputs[0]
        # Concatenate user-level features for head 1
        harm_flag_input = torch.cat((sequence_output, age.view(-1, 1), prev_flags.view(-1, 1)))
        # Feed the RoBERTa embeddings + user-level features to head 1
        logits_harm_flag = self.harm_flag_classifier(harm_flag_input)
        # Concatenate head 1 output to head 2's input
        harm_type_input = torch.cat((sequence_output, logits_harm_flag))
        # Feed the RoBERTa embeddings to head 2
        logits_harm_type = self.harm_type_classifier(harm_type_input)

        # We now return a tuple with output for each head
        return logits_harm_flag, logits_harm_type

Now we’ve got some seriously fun computation graphs to wrestle with.¹⁵ But you don’t have to worry to much about that—PyTorch’s Autograd is taking care of propagating gradients across the entire nn.Module behind the scenes! Gradients are not just being passed from head to backbone—each head can now learn jointly alongside each other as well.

If that gets you wondering how we train hydranets, then stay tuned for the next blog post in this series!

References