Auswahl der besten Modellgröße und Datensatzgröße unter einem festen Funds für LLMs

Einführung

Bei der Entwicklung von Sprachmodellen (LLMs) sind wir ständig durch Budgets eingeschränkt. Eine solche Einschränkung führt zu einem grundlegenden Kompromiss: Stellen Sie sich vor, wenn Sie ein Rechenbudget festlegen, bedeutet eine Erhöhung der Modellgröße, dass Sie die Modellgröße reduzieren müssen, auf der Sie trainieren können, und umgekehrt. Sie stellen additionally die Frage:

Sollten wir einem Modell mit mehr Parametern mehr zuweisen oder sollten wir es auf mehr Daten trainieren?

Insbesondere die Leistung und Effizienz von LLMs werden stark von diesem Kompromiss beeinflusst. Daher ist es entscheidend, ein optimales Gleichgewicht zwischen der Anzahl der Parameter eines Modells und der Anzahl der verwendeten Token zu finden.

Die gesamte Trainingsberechnung eines Transformators skaliert ungefähr wie folgt: C∝N×D, wobei

• N ist die Anzahl der Modellparameter.
• D ist die Anzahl der Token.
• C ist das feste Rechenbudget.

Es ist leicht zu erkennen, dass für ein festes C N und D umgekehrt proportional zueinander sind.

Frühere Studien (Kaplan et al., 2020; Hoffmann et al., 2022) haben herausgefunden, dass der Trainingsverlust von Modellen für maschinelles Lernen einem Potenzgesetz mit Berechnung folgt: L(C)∝C^{−α} und die optimum Modellgröße und Datensatzgrößenskala mit Berechnung als: N_opt∝C^a, D_opt∝C^b für einige constructive Werte a und b.

In diesem Artikel werden wir winzige Transformer verwenden, um zu untersuchen, wie man N und D unter einem festen Rechenwert C ausgleicht.

Versuchsaufbau

Wir entwerfen ein minimales Transformatormodell und nennen es „winziger Transformator“ mit den folgenden konfigurierbaren Eigenschaften, die die Parametergröße des Modells beeinflussen:

• Modelldimension (d_model)
• MLP-Dimension (d_mlp)
• Anzahl der Schichten (n_layers​)

Wir möchten den Transformator verschiedener Konfigurationen auf tokenisierten Sequenzen der Länge 64 des WikiText-2-Datensatzes trainieren.

Um den Effekt der Skalierung zu untersuchen, haben wir ein Raster von Modellen von sehr klein (16 versteckte Einheiten, 1 Ebene) bis relativ groß (128 versteckte Einheiten, 4 Ebenen) definiert und sie mit einer Reihe von Token von 5.000 bis 1 Mio. kombiniert. Siehe den folgenden Code:

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},   
)
# variety of tokens (D) we prepare on — simulated by way of few steps × batch × seq_len
token_budgets = (5e3, 1e4, 3e4, 5e4, 1e5, 3e5, 5e5, 1e6)  # small for demo

Indem wir den Rechenaufwand als C≈N×D annähern, besteht unsere Idee darin, die Verlustfunktion für jedes (N,D)-Paar zu berechnen und das Paar (N,D) zu finden, mit dem das Modell die minimale Verlustfunktion für ein gegebenes C erreicht: Dies ist das Gleichgewicht, nach dem wir suchen.

Umsetzung und Beobachtungen

Wir verwenden den folgenden Code, um das Modell bis zu einer festen Anzahl von Schritten mit unterschiedlichen (N,D)-Paaren zu trainieren und das Ergebnis aufzuzeichnen.

outcomes = ()
gadget = "cuda" if torch.cuda.is_available() else "cpu"

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

Dann tragen wir den endgültigen Verlust gegen die Berechnung (N×D) auf:

Wir haben folgende wichtige Beobachtungen:

1. Für kleine Rechenbudgetsschneiden kleine Modelle, die mit den meisten verfügbaren Daten trainiert wurden, besser ab als größere Modelle, die mit sehr wenigen Daten trainiert wurden.
2. Für große RechenbudgetsGrößere Modelle werden besser, wenn genügend Daten verfügbar sind.
3. Die optimale Modellgröße wächst nicht linear mit dem Rechenbudget. Beispielsweise führt eine Verdoppelung der Rechenleistung nicht wirklich zu einer doppelt so optimalen Parameteranzahl wie zuvor.

Das folgende Diagramm zeigt die Effizienzgrenze über die Modellgröße hinweg, d. h. die Menge der Modellgrößen, die für eine bestimmte Berechnung den geringsten Verlust aufweisen.

„Bestes“ Modell

Um das „beste“ Modell zu ermitteln, würden wir das Paar aus Modellgröße und Anzahl der Token auswählen, das den Verlust bei einem festen Funds minimiert.

Wir gehen davon aus, dass beide einer Potenzgesetzbeziehung folgen: N_opt∝C^α, D_opt∝C^β, und wir möchten die unbekannten Exponenten α und β durch die folgenden Schritte schätzen:

1. Nehmen Sie den Logarithmus der Größen: log?(N_opt)=αlog?(C)+const, log?(D_opt)=βlog?(C)+const.
2. Passen Sie eine lineare Regression an. Die Steigung der Regression ist nichts anderes als der Exponent des Potenzgesetzes.

Der folgende Code gibt eine solche Regression:

# Match 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))

In unserem Spielzeugexperiment haben wir herausgefunden, dass N_opt ~C^0,14 und D_opt~ C^0,86. Dieses Ergebnis zeigt möglicherweise nicht das gesamte Bild, da wir das Experiment mit vereinfachten Modellen und Konfigurationen durchgeführt haben. Wir können jedoch immer noch erkennen, dass das Wachstum der Datenverarbeitung zu einer Zunahme der optimalen Modellgröße führt, allerdings mit abnehmender Geschwindigkeit. Natürlich sollte das verbleibende Funds für mehr Schulungsgutscheine verwendet werden.

Darüber hinaus ergibt die obige Berechnung, dass das beste Verhältnis N_opt/D_opt=C^-0,72 ist. Dies bedeutet, dass Sie bei einer Erhöhung der Rechenleistung mehr Trainingstoken hinzufügen sollten, anstatt die Modellgröße zu erhöhen.

Praktische Erkenntnisse

Aus diesem Experiment, obwohl es sich um einen Spielzeugkoffer handelt, können wir mehrere Erkenntnisse gewinnen:

1. Bei einem festen Funds kann die Verwendung eines mittleren Modells mit mehr Daten die Leistung eines sehr großen Modells mit begrenzten Daten übertreffen.
2. Die optimale Modellgröße und Datengröße wachsen mit der Rechenleistung. Trainieren Sie kein Modell mit vielen Parametern, wenn Sie über ein kleines Funds verfügen.
3. Wenn das Funds steigt, berücksichtigen Sie zunächst das optimale Verhältnis N_opt/D_opt, um zu bestimmen, ob Sie die Modellgröße erhöhen oder mehr Trainingsdaten hinzufügen sollten.

Abschluss

In diesem Blogbeitrag stellen wir eine Studie zum Kompromiss zwischen Modellgröße und Daten bei einem festen Rechenbudget für LLMs mit einem Spielzeugkoffer bereit. Das Experiment zeigt, dass wir das optimale Paar aus Modellgröße und Token-Anzahl finden können, um mit einem gegebenen Funds die beste Modellleistung zu erzielen, was es Forschern und Praktikern ermöglicht, LLMs sinnvoll zu entwerfen und die besten Ergebnisse zu erzielen.

Referenz

(1) Kaplan, J., McCandlish, S., Henighan, T., Brown, TB, Chess, B., Baby, R., Grey, S., Radford, A., Wu, J. und Amodei, D. (2020). Skalierungsgesetze für neuronale Sprachmodelle.

(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, Okay., van den Driessche, G., Damoc, B., Man, A., Osindero, S., Simonyan, Okay., Elsen, E., … Sifre, L. (2022). Coaching rechenoptimaler großer Sprachmodelle.