Method and apparatus for fine-tuning machine learning models

Abstract

A computer-implemented method of fine-tuning a pre-trained machine learning (ML) model for a downstream task. The method uses the pre-trained ML model, a subset of a training dataset and scale variables. Each scale variable is associated with a parameter of the model and indicates an importance of a difference between a value of a parameter of the pre-trained ML model and a value of a parameter when the pre-trained ML model is optimised for the downstream task. The method then uses the scale variables to identify a set of updatable parameters and a set of non-updatable (“frozen”) parameters and fine-tunes the model using the training dataset by updating only the updatable parameters. The scale variables may be calculated by using stochastic gradient Langevin dynamics while minimising a loss function that includes a loss associated with the downstream task and a penalty that penalises small deviations of values of the parameters of the fine-tuned ML model from values of the parameters of the pretrained ML model. The pre-trained model may be a foundation model such as a pre-trained Vision Transformer (ViT) for processing images or a pre-trained Large Language Model, LLM for processing text.

Images