Question about the shape of my lr-find curve

[linzhang-allan]

I used the lr-finder to find the steepest gradient, but the shape of the curve is completely different from the example curve. Is this normal?

[NaleRaphael]

Can you share the code or explain what you was trying to run with LRFinder? Otherwise there is no sufficient information for me to help.

My initial guess is that you might try running LRFinder in a training loop without resetting its state by calling lr_finder.reset(), so that things just blow up.

Another suggestion is: Try to run your training script without LRFinder first. Make sure the training result is actually converging without problems, then use LRFinder to help you tune the learning rate. Don't stack too many things together at once if you are not really familiar with how those components work, it would make it difficult to debug.

As always, feel free to let me know if you need further help.

[linzhang-allan]

Below is my main code to get this curve. I just want to find the best learning rate for my initial model training. However, I plan to use the ExponentialLR scheduler to gradually decrease the learning rate as the training progresses.

def main():
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    esm_model, alphabet = pretrained.esm2_t33_650M_UR50D()
    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = SequenceInteractionModel(esm_model)
    model.to(device)
    if args.train_begin == 0: 
        train_dataset, test_dataset, start_epoch = train_begin(args, alphabet)
        optimizer = optim.Adam(model.parameters(),lr=args.lr, betas=(0.9, 0.98), eps=1e-8, weight_decay = 0.01)
    else :
        train_dataset, test_dataset, model, optimizer, start_epoch = train_last(args, model, device)
    if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = nn.DataParallel(model)   

    criterion = nn.MSELoss()
    train_loader = create_dataloader(args, train_dataset)

    lr_finder = LRFinder(model, optimizer, criterion, device="cuda")
    lr_finder.range_test(train_loader, end_lr=0.1, num_iter=1000)
    
    lrs = np.array(lr_finder.history["lr"])
    losses = np.array(lr_finder.history["loss"])

    min_grad_idx = None
    try:
        min_grad_idx = (np.gradient(np.array(losses))).argmin()
    except ValueError:
        print("Failed to compute the gradients, there might not be enough points.")
    if min_grad_idx is not None:
        best_lr = lrs[min_grad_idx]
    else:
        print("None")

    print(f"Best lr log:", best_lr)
    lr_finder.plot()
    plt.gcf().savefig("lr_finder_plot_log.png", dpi=300, bbox_inches='tight') 
    plt.close() 
    lr_finder.reset()

[NaleRaphael]

Sorry, I have a limited understanding to protein language model, so please correct me if there is anything wrong about it.

According to your training script, you are using esm2_t33_650M_UR50D as a base model for sequence interaction modeling. And based on things and information I can find for now, it usually uses MSA (Multiple Sequence Alignment)-related loss, which seems very different to MSELoss. So, can you check whether it's actually a working loss function for your case?

Otherwise, you might need to run the training without LRFinder first then check whether the result could converge, then rerun the training with LRFinder again to tune the model with a better performance.

[linzhang-allan]

Thank you for your response. My objective is to employ a sequence interaction model for a regression task, specifically to predict traits of various proteins. Since I've been unable to identify a learning rate that allows my model to converge, I'm questioning whether the issue lies with my data or the learning rate itself. Interestingly, when I switched to a different dataset, I obtained a relatively normal convergence curve. This leads me to believe that using lr-find might also be an effective tool to evaluate the quality of my data. What is your understanding？

[NaleRaphael]

Sounds interesting, and I get your point now, because I was assuming you were trying to tune the model.

For this kind of problem, I would try to use a fixed learning rate to run through the dataset first (just guess a "reasonable" learning rate for it though it might be not that easy to guess), then check the convergence. Usually, model should be able to converge even it's not actually learning.

Otherwise, I would start investigating the loss calculated for each mini-batch of data. Like inserting pdb.set_trace() to the training loop and print the loss value then check it with the data in that batch. But I'm not sure whether this is possible or easy for your dataset.

If you want to check whether it's a data quality problem, maybe you can try to split the dataset into small parts, then run the training script with each small parts one by one to pinpoint the problem.

These are what I can suggest for now, try the easier approach first.
Hope you can solve this problem in the end, and feel free to let me know if you need help.

[NaleRaphael]

A few other things I just came up with:

1. Sometimes the model needs to run through quite a few batches to start learning something from data, then the loss + other metrics would start converging. And based on the figure in OP, I noticed that the number of data points shown in it doesn't seem to match the ~1000 points suggested by the code snippet you provided below (num_iter=1000).
   Not sure if it's because the training was terminated earlier or due to other reason, but the figure shouldn't look like this if everything is running normally. It appears to show <100 data points, even though there would be 15 data points truncated by lr_finder.plot() with default settings (see also here). So this might be a thing to check.

2. Also, the loss range is around 0.7 to 1.1. While we cannot say whether it's a large loss for the dataset, I would suspect it's a kind of vibration resulted by the cold start for the downstream layers. In this case, I would suggest running the model for more iterations (with a fixed learning rate of course) and checking the loss history to see how it evolves.

But if the figure is not the exact one generated by the code snippet (something has been modified in these days), you can safely ignore 1.

[linzhang-allan]

Thank you for your suggestions. I'll definitely give them a try.