Add cosine annealing LR scheduler to ChemPropLightningModule

[smcolby]

Problem

ChemPropLightningModule uses a fixed learning rate throughout training. A fixed LR often leads to loss oscillation in late epochs and leaves
performance gains on the table that LR decay would capture.

Proposed Solution

Add CosineAnnealingLR as an optional scheduler in configure_optimizers():

# moal/model.py
from torch.optim.lr_scheduler import CosineAnnealingLR

def configure_optimizers(self):
    optimizer = ...  # existing AdamW setup
    if self.use_lr_scheduler:
        scheduler = CosineAnnealingLR(optimizer, T_max=self.max_epochs, eta_min=1e-7)
        return {
            "optimizer": optimizer,
            "lr_scheduler": {"scheduler": scheduler, "interval": "epoch"},
        }
    return optimizer

New constructor parameters:

use_lr_scheduler: bool = True

max_epochs is already available via the Lightning Trainer but needs to be stored at construction time (or passed to refit()).

Files

• moal/model.py

Notes

For the freeze/unfreeze schedule (freeze_epochs), the scheduler should reset or use a separate warm-up phase to avoid decaying the LR before the encoder is unfrozen.