What is noise-rater?

NoiseRater is a Python framework for improving how diffusion models are trained. Diffusion models are a class of AI image-generation model that work by learning to reverse a process of gradually adding random noise to an image. In standard training, the noise added at each step is drawn randomly. NoiseRater's insight is that not all noise samples are equally useful for learning, some help the model improve faster than others, so it trains a small scoring function, called a noise rater, to select better noise at each training step. The scoring function is trained using a technique called bilevel optimization. This involves two nested training loops: an inner loop that trains the main image model using a diffusion loss weighted by the noise rater's scores, and an outer loop that evaluates that inner model's quality on a separate validation set and sends feedback back through the inner loop to update the noise rater. The result is a learned policy for picking noise that improves convergence of the underlying model. The pipeline has three stages: first train a baseline model to obtain reference checkpoints, then train the noise rater using those checkpoints as the inner model, then resume or restart main model training with the learned rater guiding noise selection. Training is distributed across multiple GPUs using torchrun. The framework expects image data in ImageNet's ImageFolder directory format.