Choisir la meilleure taille de modèle et la meilleure taille d’ensemble de données dans le cadre d’un budget fixe pour les LLM

Introduction

modèles de langage (LLM), nous sommes toujours contraints par les budgets. Une telle contrainte conduit à un compromis fondamental : imaginez que si vous fixez un budget de calcul, augmenter la taille du modèle signifie que vous devez réduire la taille du modèle sur lequel vous pouvez vous entraîner, et vice versa. Alors vous posez la question :

Devrions-nous allouer davantage à un modèle avec plus de paramètres, ou devrions-nous l’entraîner sur plus de données ?

En particulier, les performances et l’efficacité des LLM sont largement influencées par ce compromis. Il est donc crucial de trouver un équilibre optimal entre le nombre de paramètres d’un modèle et le nombre de tokens utilisés.

Le calcul d’entraînement total d’un transformateur s’échelonne approximativement comme : C∝N×D, où

• N est le nombre de paramètres du modèle.
• D est le nombre de jetons.
• C est le budget de calcul fixe.

Il est simple de voir que pour un C fixe, N et D sont inversement proportionnels l’un à l’autre.

Des études antérieures (Kaplan et al., 2020 ; Hoffmann et al., 2022) ont montré que la perte d’entraînement des modèles d’apprentissage automatique suit une loi de puissance avec calcul : L(C)∝C^{−α} et le optimal taille du modèle et échelle de taille de l’ensemble de données avec calcul comme : N_opt∝C^a, D_opt∝C^b pour certaines valeurs positives a et b.

Dans cet article, nous utiliserons de minuscules Transformers pour explorer comment équilibrer N et D sous un calcul C fixe.

Configuration de l’expérience

Nous concevons un modèle de transformateur minimal, et nous l’appelons « petit transformateur » avec les propriétés configurables suivantes qui influencent la taille des paramètres du modèle :

• Dimension du modèle (d_model)
• Dimension MLP (d_mlp)
• Nombre de couches (n_layers​)

Nous aimerions entraîner le transformateur de différentes configurations sur des séquences tokenisées de longueur 64 de l’ensemble de données WikiText-2.

Pour étudier l’effet de la mise à l’échelle, nous avons défini une grille de modèles allant du très petit (16 unités cachées, 1 couche) au relativement grand (128 unités cachées, 4 couches) et les avons combinés avec une gamme de jetons de 5 000 à 1 M. Voir le code ci-dessous :

model_configs = [
    {"d_model": 16,  "d_mlp": 64,   "n_layers": 1},  
    {"d_model": 24,  "d_mlp": 96,   "n_layers": 1},   
    {"d_model": 32,  "d_mlp": 128,  "n_layers": 2},
    {"d_model": 48,  "d_mlp": 192,  "n_layers": 2},
    {"d_model": 64,  "d_mlp": 256,  "n_layers": 3},
    {"d_model": 96,  "d_mlp": 384,  "n_layers": 3},
    {"d_model": 128, "d_mlp": 512,  "n_layers": 4},   
]
# number of tokens (D) we train on — simulated via few steps × batch × seq_len
token_budgets = [5e3, 1e4, 3e4, 5e4, 1e5, 3e5, 5e5, 1e6]  # small for demo

En approximant le coût de calcul à C≈N×D, notre idée est de calculer la fonction de perte pour chaque paire (N,D) et de trouver la paire (N,D) avec laquelle le modèle atteint la fonction de perte minimale pour un C donné : c’est l’équilibre que nous recherchons.

Mise en œuvre et observations

Nous utilisons le code ci-dessous pour entraîner le modèle jusqu’à un nombre fixe d’étapes avec différentes paires (N, D) et enregistrer le résultat.

results = []
device = "cuda" if torch.cuda.is_available() else "cpu"

for cfg in model_configs:
    model = TinyTransformer(vocab_size=len(tokenizer), **cfg)
    N_params = count_params(model)
    for D in token_budgets:
        steps = int(D // (SEQ_LEN * 16))  # assuming batch_size=16
        dataloader = DataLoader(
            tokenized_dataset["train"].shuffle(seed=0),
            batch_size=16,
            collate_fn=collate_fn
        )
        avg_loss = train_one(model, dataloader, steps=steps, device=device)
        compute = N_params * D
        results.append({
            "N": N_params,
            "D": D,
            "C": compute,
            "loss": avg_loss
        })

Nous traçons ensuite la perte finale par rapport au calcul (N×D) :

Nous avons les observations importantes suivantes :

1. Pour les petits budgets de calculles petits modèles entraînés sur la plupart des données disponibles fonctionnent mieux que les modèles plus grands entraînés sur très peu de données.
2. Pour les gros budgets de calculles modèles plus grands deviennent meilleurs lorsque suffisamment de données sont disponibles.
3. La taille optimale du modèle n’augmente pas de manière linéaire avec le budget de calcul. Par exemple, doubler le calcul ne conduit pas vraiment à un nombre optimal de paramètres deux fois comme avant.

Le graphique ci-dessous donne la frontière efficace selon la taille du modèle, c’est-à-dire l’ensemble des tailles de modèle qui présentent la perte la plus faible pour un calcul donné.

« Meilleur » modèle

Pour déterminer le « meilleur » modèle, nous sélectionnerions la paire de tailles de modèle et le nombre de jetons qui minimisent les pertes avec un budget fixe.

Nous supposons que les deux suivent une relation puissance-loi : N_opt∝C^α, D_opt∝C^β, et nous aimerions estimer les exposants inconnus α et β par les étapes suivantes :

1. Prendre le logarithme des quantités : log?(N_opt)=αlog?(C)+const, log?(D_opt)=βlog?(C)+const.
2. Ajustez une régression linéaire. La pente de la régression n’est rien d’autre que l’exposant de la loi de puissance.

Le code suivant donne une telle régression :

# Fit log-log linear regression
a_slope, a_intercept, *_ = st.linregress(np.log(frontier.C), np.log(frontier.N))
b_slope, b_intercept, *_ = st.linregress(np.log(frontier.C), np.log(frontier.D))

Dans notre expérience sur les jouets, nous avons constaté que N_opt ~C^0,14 et D_opt~ C^0,86. Ce résultat pourrait ne pas révéler l’image entière car nous avons fait l’expérience sur un modèle et des configurations simplifiés. Mais nous pouvons toujours constater que la croissance de l’informatique entraîne une augmentation de la taille optimale des modèles, mais à un rythme décroissant. De toute évidence, le budget restant devrait être attribué à davantage de jetons de formation.

De plus, le calcul ci-dessus donne le fait que le meilleur rapport N_opt/D_opt=C^-0,72. Cela implique que lorsque vous augmentez le calcul, vous devez ajouter plus de jetons de formation plutôt que d’augmenter la taille du modèle.

Points pratiques à retenir

De cette expérience, bien qu’il s’agisse d’un cas de jouet, nous pouvons extraire plusieurs enseignements :

1. Pour un budget fixe, l’utilisation d’un modèle moyen avec plus de données peut surpasser un très grand modèle avec des données limitées.
2. La taille optimale du modèle et la taille des données augmentent avec le calcul. N’entraînez pas un modèle avec de nombreux paramètres si vous disposez d’un petit budget.
3. Lorsque le budget augmente, considérez d’abord le rapport optimal N_opt/D_opt pour déterminer si vous devez augmenter la taille du modèle ou ajouter plus de données d’entraînement.

Conclusion

Dans cet article de blog, nous proposons une étude du compromis entre la taille du modèle et les données dans le cadre d’un budget de calcul fixe pour les LLM avec une mallette de jouets. L’expérience montre que nous pouvons trouver la paire optimale de taille de modèle et de nombre de jetons pour obtenir les meilleures performances de modèle avec un budget donné, permettant ainsi aux chercheurs et aux praticiens de concevoir judicieusement des LLM et d’obtenir les meilleurs résultats.

Référence

[1] Kaplan, J., McCandlish, S., Henighan, T., Brown, TB, Chess, B., Child, R., Gray, S., Radford, A., Wu, J. et Amodei, D. (2020). Lois de mise à l’échelle pour les modèles de langage neuronal.

[2] Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, LA, Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E.,… Sifre, L. (2022). Formation de grands modèles de langage optimisés pour le calcul.