Transfer Learning Distillation Model Loss Not Decreasing

Question

Currently I'm trying to reproduce paper "A Deep Transfer Learning Method for Cross-Lingual Natural Language Inference" (Bandyopadhyay et al., LREC 2022) for Cross-Lingual Natural Language Inference task. But, the model I'm trying to reproduce is not learning any parameters which demonstrated by the model's loss not decreasing.

The dataset I'm using is IndoNLI with hypothesis sentences translated into Javanese. But, as you can read on the paper, you can also use XNLI for this task.

For this experiment, I'm using Pytorch, Huggingface Transformers, Pandas, Numpy, and Wandb for logging.

First, I construct my dataset as follows:

class CompDataset(Dataset):
    def __init__(self, df_teacher, df_student):
        self.df_data_teacher = df_teacher
        self.df_data_student = df_student
        
    def __getitem__(self, index):
        # Teacher
        sentence_teacher_1 = self.df_data_teacher.loc[index, 'premise']
        sentence_teacher_2 = self.df_data_teacher.loc[index, 'hypothesis']
        
        encoded_dict_teacher = tokenizer.encode_plus(
            sentence_teacher_1,
            sentence_teacher_2,
            add_special_tokens = True,
            max_length = MAX_LEN,
            truncation='longest_first',
            padding = 'max_length',
            return_attention_mask = True,
            return_tensors = 'pt'
        )
        
        padded_token_list_teacher = encoded_dict_teacher['input_ids'][0]
        att_mask_teacher = encoded_dict_teacher['attention_mask'][0]
        tok_type_id_teacher = encoded_dict_teacher['token_type_ids'][0]
        
        target_teacher = torch.tensor([self.df_data_teacher.loc[index, 'label']])
        lt_target_teacher = torch.LongTensor(target_teacher)
        onehot_encoded_lbl_teacher = F.one_hot(lt_target_teacher, num_classes=3) # 3 classes: entails, neutral, contradict
        
        # Student
        sentence_student_1 = self.df_data_student.loc[index, 'premise']
        sentence_student_2 = self.df_data_student.loc[index, 'hypothesis']
        
        encoded_dict_student = tokenizer.encode_plus(
            sentence_student_1,
            sentence_student_2,
            add_special_tokens = True,
            max_length = MAX_LEN,
            truncation='longest_first',
            padding = 'max_length',
            return_attention_mask = True,
            return_tensors = 'pt'
        )
        
        padded_token_list_student = encoded_dict_student['input_ids'][0]
        att_mask_student = encoded_dict_student['attention_mask'][0]
        tok_type_id_student = encoded_dict_student['token_type_ids'][0]
        
        target_student = torch.tensor([self.df_data_student.loc[index, 'label']])
        lt_target_student = torch.LongTensor(target_student)
        onehot_encoded_lbl_student = F.one_hot(lt_target_student, num_classes=3) # 3 classes: entails, neutral, contradict
        
        output = {
            "input_ids_teacher": padded_token_list_teacher, 
            "attention_mask_teacher": att_mask_teacher,
            "token_type_ids_teacher": tok_type_id_teacher,
            "lbl_teacher": onehot_encoded_lbl_teacher,
            "input_ids_student": padded_token_list_student, 
            "attention_mask_student": att_mask_student,
            "token_type_ids_student": tok_type_id_student,
            "lbl_student": onehot_encoded_lbl_student
        }
        
        return output
    
    def __len__(self):
        return len(self.df_data_teacher)

Then, I build the transformers' dataset & dataloader. The df_train_t and df_train_student being dataframe for teacher dataset (Indonesian premise-Indonesian hypothesis) and student dataset (Indonesian premise-Javanese hypothesis).

train_data_cmp = CompDataset(df_train_t, df_train_student)
valid_data_cmp = CompDataset(df_valid_t, df_valid_student)
test_data_cmp = CompDataset(df_test_t, df_test_student)

train_dataloader = DataLoader(train_data_cmp, batch_size = BATCH_SIZE)
valid_dataloader = DataLoader(valid_data_cmp, batch_size = BATCH_SIZE)
test_dataloader = DataLoader(test_data_cmp, batch_size = BATCH_SIZE)

After that, I try to build the model using the schematic and algorithm of transfer learning method provided on the paper. As you can see on the code below, I tried to freeze the mBERT model for teacher, and update only the student model parameters.

class TransferLearningPaper(PreTrainedModel):
    def __init__(self, config, lambda_kld, learningrate_student, batchnorm_epsilon = 1e-5):
        super(TransferLearningPaper, self).__init__(config)
        
        self.bert_model_teacher = BertModel.from_pretrained(
            MODEL_TEACHER_TYPE, # using already pretrained mBERT in INA language
            num_labels = 3,
            output_hidden_states=True
        )
        
        # Freeze teacher mBERT parameters
        for params_teacher in self.bert_model_teacher.parameters():
            params_teacher.requires_grad = False
    
        self.bert_model_student = BertModel.from_pretrained(
            MBERT_TYPE,
            num_labels = 3,
            output_hidden_states=True
        )
        
        self.optimizer_student = AdamW(
            self.bert_model_student.parameters(), 
            lr=learningrate_student
        )
        
        self.linear = nn.Linear(config.hidden_size, 3)  # Linear layer
        self.batchnorm = nn.BatchNorm1d(config.hidden_size, eps=batchnorm_epsilon)
        self.softmax = nn.Softmax(dim=1)  # Softmax activation
        
        self.cross_entropy = nn.CrossEntropyLoss()
        self.kld = nn.KLDivLoss(reduction='batchmean')
        
        # Initialize the weights of the linear layer
        self.linear.weight.data.normal_(mean=0.0, std=0.02)
        self.linear.bias.data.zero_()
        
        self.lambda_kld = lambda_kld
    
    def forward(self, input_ids_teacher, attention_mask_teacher, token_type_ids_teacher, lbl_teacher, input_ids_student, attention_mask_student, token_type_ids_student, lbl_student):
        # assume the label is already one-hot encoded
        
        self.bert_model_teacher.eval()
        self.bert_model_student.eval()
        
        with torch.no_grad():
            outputs_teacher = self.bert_model_teacher(
                input_ids=input_ids_teacher, 
                attention_mask=attention_mask_teacher, 
                token_type_ids=token_type_ids_teacher
            )
            outputs_student = self.bert_model_student(
                input_ids=input_ids_student, 
                attention_mask=attention_mask_student, 
                token_type_ids=token_type_ids_student
            )
        
            # take CLS token of the last hidden state
            pooled_output_teacher = outputs_teacher[0][:, 0, :]
            pooled_output_student = outputs_student[0][:, 0, :]
        
        batchnormed_logits = self.batchnorm(pooled_output_student)
        linear_output = self.linear(batchnormed_logits) # the output's logits
        softmax_linear_output = F.log_softmax(linear_output, dim=1)
        
        lbl_student = lbl_student[:,0,:].float()
        lbl_teacher = lbl_teacher[:,0,:].float()
        softmax_linear_output = softmax_linear_output.float()
        
        cross_entropy_loss = self.cross_entropy(softmax_linear_output, lbl_student)
        total_kld = self.kld(F.log_softmax(pooled_output_student, dim=1), F.softmax(pooled_output_teacher, dim=1))
        
        joint_loss = cross_entropy_loss + (self.lambda_kld * total_kld )
        
        return {"loss": joint_loss, "logits": softmax_linear_output}
    
    def update_param_student_model(self, loss):
        # Doing customized backpropagation for student's model
        self.bert_model_student.train()
        
        self.optimizer_student.zero_grad()
        loss.backward()
        self.optimizer_student.step()

Then, I instantiate the model and its configurations and hyperparameters:

config = PretrainedConfig(
    problem_type = "single_label_classification",
    id2label = {
        "0": "ENTAIL",
        "1": "NEUTRAL",
        "2": "CONTRADICTION"
    },
    label2id = {
        "ENTAIL": 0,
        "NEUTRAL": 1,
        "CONTRADICTION": 2
    },
    num_labels = 3,
    hidden_size = 768,
    name_or_path = "indojavanesenli-transfer-learning",
    finetuning_task = "indonesian-javanese natural language inference"
)
print(config)
transferlearning_model = TransferLearningPaper(
    config = config,
    lambda_kld = 0.011, # antara 0.01-0.5
    learningrate_student = STUDENT_LRATE,
    batchnorm_epsilon = BATCH_NORM_EPSILON
)
transferlearning_model = transferlearning_model.to(device)

After that, I create functions to train and validate my model:

def train(the_model, train_data):
    the_model.train()
    
    batch_loss = 0
    
    for batch, data in enumerate(train_data):
        input_ids_teacher = data["input_ids_teacher"].to(device)
        attention_mask_teacher = data["attention_mask_teacher"].to(device)
        token_type_ids_teacher = data["token_type_ids_teacher"].to(device)
        lbl_teacher = data["lbl_teacher"].to(device)
        input_ids_student = data["input_ids_student"].to(device)
        attention_mask_student = data["attention_mask_student"].to(device)
        token_type_ids_student = data["token_type_ids_student"].to(device)
        lbl_student = data["lbl_student"].to(device)
        
        output = the_model(
            input_ids_teacher = input_ids_teacher, 
            attention_mask_teacher = attention_mask_teacher, 
            token_type_ids_teacher = token_type_ids_teacher, 
            lbl_teacher = lbl_teacher, 
            input_ids_student = input_ids_student, 
            attention_mask_student = attention_mask_student, 
            token_type_ids_student = token_type_ids_student, 
            lbl_student = lbl_student
        )
        
        loss_model = output["loss"]
        batch_loss += loss_model
        wandb.log({"train/loss": loss_model})
        
        # Backpropagation
        the_model.update_param_student_model(loss_model)
    
    training_loss = batch_loss / BATCH_SIZE
    
    return training_loss

def validate(the_model, valid_data):
    the_model.eval()
    
    batch_loss = 0
    
    with torch.no_grad():
        for batch, data in enumerate(valid_data):
            input_ids_teacher = data["input_ids_teacher"].to(device)
            attention_mask_teacher = data["attention_mask_teacher"].to(device)
            token_type_ids_teacher = data["token_type_ids_teacher"].to(device)
            lbl_teacher = data["lbl_teacher"].to(device)
            input_ids_student = data["input_ids_student"].to(device)
            attention_mask_student = data["attention_mask_student"].to(device)
            token_type_ids_student = data["token_type_ids_student"].to(device)
            lbl_student = data["lbl_student"].to(device)

            output = the_model(
                input_ids_teacher = input_ids_teacher, 
                attention_mask_teacher = attention_mask_teacher, 
                token_type_ids_teacher = token_type_ids_teacher, 
                lbl_teacher = lbl_teacher, 
                input_ids_student = input_ids_student, 
                attention_mask_student = attention_mask_student, 
                token_type_ids_student = token_type_ids_student, 
                lbl_student = lbl_student
            )

            logits = output["logits"].cpu().detach().numpy()
            packed_val = logits, lbl_student.cpu().detach().numpy()
            metrics = compute_metrics(packed_val)

            loss_model = output["loss"]
            batch_loss += loss_model
            wandb.log({
                "eval/loss": loss_model, 
                "eval/f1_score": metrics["f1_score"], 
                "eval/accuracy": metrics["accuracy"],
                "eval/precision": metrics["precision"],
                "eval/recall": metrics["recall"]
            })
    
        eval_loss = batch_loss / BATCH_SIZE
    
    return eval_loss, metrics

def training_sequence(the_model, train_data, valid_data, epochs):
    track_train_loss = []
    track_val_loss = []
    
    t = trange(epochs, colour="green", position=0, leave=True)
    for ep in t:
        training_loss = train(the_model, train_data)
        valid_loss, _ = validate(the_model, valid_data)
        
        track_train_loss.append(training_loss)
        track_val_loss.append(valid_loss)
        
        t.set_description(f"Epoch [{ep + 1}/{epochs}] - Train loss: {training_loss:.2f} Valid loss: {valid_loss:.2f}")
        
        if valid_loss < min(track_val_loss) or ep + 1 == 1:
            the_model.save_pretrained(
                save_directory = MODEL_PATH + "indojavanesenli-transfer-learning"
            )
            
        wandb.log({
            "train_loss/epoch": training_loss,
            "validation_loss/epoch": valid_loss
        })
        
    return {
        "training_loss": track_train_loss,
        "validation_loss": track_val_loss
    }

Finally, I train my model by using:

training_result = training_sequence(transferlearning_model, train_dataloader, valid_dataloader, NUM_EPOCHS)

But the problem is, during training, the model not updating the student's model parameters as you can see on Fig.1 below.

Figure 1. Model loss not decreasing

FYI, this is the configuration variable I use for the code above:

TOKENIZER_TYPE = 'bert-base-multilingual-cased'
MBERT_TYPE = 'bert-base-multilingual-cased'
MODEL_TEACHER_TYPE = 'jalaluddin94/nli_mbert' # This is an already fine-tuned mBERT on the Indonesian language
MODEL_PATH = 'D:/Training/Machine Learning/NLP/NLI/Indo-Javanese-NLI/ResearchedModels/'

STUDENT_LRATE = 2e-5
MAX_LEN = 512
NUM_EPOCHS = 25
BATCH_SIZE = 12
BATCH_NORM_EPSILON = 1e-5
LAMBDA_L2 = 3e-5

Answer 1

I have figured it out.

It's because of with torch.no_grad(): on the forward pass of the model.

The BERT student parameters wouldn't be updated because it lost its gradient.