Hello,
There are times when I need to fine-tune the output head of the pre-trained Transformer models, for instance, the multiple-choice head of the GPT2DoubleHeadsModel is not pre-trained, so I need to fine-tune the multiple-choice head, although the weights for the main body of the model are okay. People are telling me that I am supposed to set my learning rate fairly high for the weights of the output head, while setting them low for the weights of the main body, but I don’t know how to do that.
The way I train the pre-trained transformer models is like below:
def train(model, train_data, optimizer, scheduler, log_interval, pad_index):

    # turn on a training mode
    model.train()
    
    # initialize total_loss to 0
    total_loss = 0

    for batch_index, batch in enumerate(train_data):
        
        gc.collect()
        
        input_ids = [instance for instance in batch.text]

        optimizer.zero_grad() 
        
        input_ids = torch.tensor([input_ids], dtype=torch.long)
        
        loss = model(input_ids, labels = input_ids)[0]
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        
        optimizer.step()
        
        total_loss = total_loss + loss 

        if batch_index % log_interval == 0 and batch_index > 0:
            cur_loss = total_loss / log_interval
            print('| epoch {:3d} | {:5d}/{:5d} batches | lr {:02.9f} | loss {:5.4f} | ppl {:8.4f}'.format(
                    epoch, batch_index, len(train_iter), scheduler.get_lr()[0], cur_loss, math.exp(cur_loss)))
                   
            total_loss = 0 
        
        del input_ids
        del loss
        gc.collect() 


optimizer = AdamW(model.parameters(), lr = 0.000005200, correct_bias = True)
scheduler = get_constant_schedule(optimizer = optimizer, last_epoch = -1)

for epoch in range(1, epochs + 1):
    
    gc.collect()
    
    epoch_start_time = time.time()
    
    train(model, train_data, 
                  optimizer, scheduler, 
                  log_int, pad_index)
    
    print('-' * 89)
    print('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.4f} | '
          'valid ppl {:8.4f}'.format(epoch, (time.time() - epoch_start_time),
                                     val_loss, math.exp(val_loss)))
    print('-' * 89)

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_model = model
 
    scheduler.step() 

What is the better way for fine-tuning the output-head?
Thank you,