AWS Strands AI: Training a 2B Model.

AI models are growing exponentially while the importance of size being directly correlated to performance is decreasing, but this does not always mean that bigger is better for every application

Smaller models can also perform at remarkable levels, if trained and evaluated accurately.

While researching AWS Strands / AI model development I stumbled onto an interesting question:

Can a 2B-class model compete with significantly larger models?

Let's dive into the challenge

A 2B parameter mark is substantially lower than most of the current top-tier language models; this, in turn, makes such models interesting from an engineering standpoint: smaller models can be quicker and more efficient than their larger counterparts

The key is to not train one model and measure its accuracy - a comprehensive performance evaluation should consider various tasks, inputs, and domains, including text and image inputs, and the ability to generalize across various tasks, including those the model was not trained on specifically.

One interesting way to tackle this issue is to utilize knowledge distillation

The idea is quite simple and intuitive:

Large Teacher Model -> Training Signals -> Smaller Student Model

The general concept is to use the information learned by a large model to train another, smaller model. Such an approach allows the smaller model to learn from the larger one and thus incorporate some of the larger model's properties, such as accuracy, without having an exceptionally large number of parameters itself.

The simplest way to evaluate such an approach is to train multiple models using various combinations of the techniques mentioned above and compare their performance using a standardized set of measurements, such as:

This approach is similar to a bake-off, where multiple potential options are trained and measured using a similar set of rules and then compared using a standardized set of metrics.

While the calibration of the final model is important (it is quite common for models to know when they are wrong), it is even more important to have a correctly calibrated teacher model for knowledge distillation.

After all, if the teacher model is wrong about something, there is a very real chance that the student will also be wrong.

While calibration is important, it should not be the sole deciding factor when choosing between multiple models; the accuracy of the model should also be taken into account.

Another important consideration is that the model should be tested on more than just text data.

The addition of image data into the mix is interesting, as it presents an additional challenge for the model: most of the current models are text-based, and the addition of vision capabilities transforms a model from a relatively straightforward text-processing machine into a more complicated multimodal system

Such a system now has two potential weak points: vision and text processing pipelines, which means that there are now two potential areas where the model can fail.

When debugging such a system, it is also important to consider all components, not just one.

Debugging a multimodal model is not substantially different from debugging any other machine learning model: you try to identify the weak point in the system and then focus on improving it.

In practice, however, this rarely means finding the issue in model architecture: more often than not, the problems are caused by subtle inconsistencies in data formatting, preprocessing or evaluation techniques.

Fixing these issues often takes a significant amount of time and effort, but it is always worth it: working with a correctly prepared dataset and a well-prepared model is substantially less error-prone than trying to get an already existing one to perform better on the same data.

One particularly important consideration is that a good performance on a benchmark task does not necessarily mean that the model will generalize well.

A more interesting approach would be to test the model on a set of tasks the model was not explicitly trained for; if the model achieves good results on such a task, that is a sign that the model has learned to perform well on the distribution of tasks that it was trained on.

Thus, it is reasonable to structure the model training and evaluation around these principles:

Lessons Learned

The most important lesson is that the size of the model is not as important as one might think: a relatively small model can achieve surprisingly good results, given the right amount of training and evaluation.

While the final accuracy of the model is important, there are other factors that should be considered when comparing different models: calibration, training data, hyperparameters, base model, etc.

The process of training and evaluating such a model can be time- and resource-consuming, but it is always rewarding: it provides insight into what makes a model good for a particular task and highlights areas that require improvement.

In the end, remember that the local model benchmarks are interesting, but they should not be the sole deciding factor when comparing different approaches.