Beyond Cosine Decay: On the effectiveness of Infinite Learning Rate Schedule for Continual Pre-training
This work addresses the challenge of adapting self-supervised learning to non-stationary data streams for AI systems, but it is incremental as it focuses on comparing and improving an existing schedule rather than introducing a new paradigm.
The paper tackled the problem of forgetting in continual pre-training with self-supervised learning by comparing the infinite learning rate schedule to the widely used repeated cosine decay, finding that the infinite schedule consistently enhances performance across image and language datasets, such as surpassing repeated cosine decay in MAE pre-training and zero-shot LM benchmarks.
The ever-growing availability of unlabeled data presents both opportunities and challenges for training artificial intelligence systems. While self-supervised learning (SSL) has emerged as a powerful paradigm for extracting meaningful representations from vast amounts of unlabeled data, existing methods still struggle to adapt to the non-stationary, non-IID nature of real-world data streams without forgetting previously learned knowledge. Recent works have adopted a repeated cosine annealing schedule for large-scale continual pre-training; however, these schedules (1) inherently cause forgetting during the re-warming phase and (2) have not been systematically compared to existing continual SSL methods. In this work, we systematically compare the widely used cosine schedule with the recently proposed infinite learning rate schedule and empirically find the latter to be a more effective alternative. Our extensive empirical evaluation across diverse image and language datasets demonstrates that the infinite learning rate schedule consistently enhances continual pre-training performance compared to a repeated cosine decay without being restricted to a fixed iteration budget. For instance, in a small-scale MAE pre-training setup, it outperforms several strong baselines from the literature. We then scale up our experiments to larger MAE pre-training and autoregressive language model pre-training. Our results show that the infinite learning rate schedule remains effective at scale, surpassing repeated cosine decay for both MAE pre-training and zero-shot LM benchmarks.