Elegir el mejor tamaño de modelo y tamaño de conjunto de datos con un presupuesto fijo para LLM

Introducción

modelos de lenguaje (LLM), siempre estamos limitados por los presupuestos. Tal restricción conduce a una compensación fundamental: imagine que si fija un presupuesto de cómputo, aumentar el tamaño del modelo significa que debe reducir el tamaño del modelo en el que puede entrenar, y viceversa. Entonces estás haciendo la pregunta:

¿Deberíamos asignar más a un modelo con más parámetros o deberíamos entrenarlo con más datos?

En particular, el desempeño y la eficiencia de los LLM se ven influenciados en gran medida por esta compensación. Por tanto, es crucial encontrar un equilibrio óptimo entre la cantidad de parámetros de un modelo y la cantidad de tokens utilizados.

El cálculo de entrenamiento total de un transformador se escala aproximadamente como: C∝N×D, donde

N es el número de parámetros del modelo. D es el número de fichas. C es el presupuesto de cálculo fijo.

Es sencillo ver que para un C fijo, N y D son inversamente proporcionales entre sí.

Estudios anteriores (Kaplan et al., 2020; Hoffmann et al., 2022) han descubierto que la pérdida de entrenamiento de los modelos de aprendizaje automático sigue una ley de potencia con cálculo: L(C)∝C^{−α} y el tamaño óptimo del modelo y la escala de tamaño del conjunto de datos con cálculo como: N_opt∝C^a, D_opt∝C^b para algunos valores positivos a y b.

En este artículo, utilizaremos pequeños transformadores para explorar cómo equilibrar N y D en un cálculo fijo C.

Configuración del experimento

Diseñamos un modelo de transformador mínimo y lo llamamos “transformador pequeño” con las siguientes propiedades configurables que influyen en el tamaño de los parámetros del modelo:

Dimensión del modelo (d_model) Dimensión MLP (d_mlp) Número de capas (n_layers​)

Nos gustaría entrenar el transformador de diferentes configuraciones en secuencias tokenizadas de longitud 64 del conjunto de datos WikiText-2.

Para estudiar el efecto del escalado, definimos una cuadrícula de modelos desde muy pequeña (16 unidades ocultas, 1 capa) hasta relativamente grande (128 unidades ocultas, 4 capas) y los combinamos con una variedad de tokens de 5k a 1M. Vea el código a continuación:

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},
]
# número de tokens (D) con los que entrenamos: simulados mediante unos pocos pasos × lote × seq_len token_budgets = [5e3, 1e4, 3e4, 5e4, 1e5, 3e5, 5e5, 1e6] # pequeño para demostración

Al aproximar el costo de cálculo como C≈N×D, nuestra idea es calcular la función de pérdida para cada par (N,D) y encontrar el par (N,D) con el que el modelo alcanza la función de pérdida mínima para un C dado: este es el equilibrio que estamos buscando.

Implementación y observaciones.

Usamos el siguiente código para entrenar el modelo hasta un número fijo de pasos con diferentes pares (N,D) y registrar el resultado.

resultados = []
dispositivo = “cuda” si torch.cuda.is_available() else “cpu” para cfg en model_configs: modelo = TinyTransformer(vocab_size=len(tokenizer), **cfg) N_params = count_params(modelo) para D en token_budgets: pasos = int(D // (SEQ_LEN * 16)) # asumiendo tamaño_lote = 16 cargador de datos = Cargador de datos (conjunto de datos_tokenizados[“train”].shuffle (semilla = 0), tamaño de lote = 16, intercalar_fn = intercalar_fn) avg_loss = train_one (modelo, cargador de datos, pasos = pasos, dispositivo = dispositivo) calcular = N_params * D resultados.append ({ “N”: N_params, “D”: D, “C”: calcular, “pérdida”: avg_loss })

Luego trazamos la pérdida final contra el cálculo (N×D):

Imagen del autor: pérdida de entrenamiento versus computación

Tenemos las siguientes observaciones importantes:

Para presupuestos de computación pequeños, los modelos pequeños entrenados con la mayoría de los datos disponibles funcionan mejor que los modelos más grandes entrenados con muy pocos datos. Para grandes presupuestos de computación, los modelos más grandes resultan mejores cuando hay suficientes datos disponibles. El tamaño óptimo del modelo no crece linealmente con el presupuesto informático. Por ejemplo, duplicar el cálculo en realidad no conduce a un número óptimo de parámetros el doble que antes.

El siguiente gráfico muestra la frontera eficiente entre los tamaños del modelo, es decir, el conjunto de tamaños de modelo que tienen la pérdida más baja para un cálculo determinado.

Imagen del autor: frontera eficiente

“Mejor” modelo

Para determinar el “mejor” modelo, seleccionaríamos el par de tamaño de modelo y la cantidad de tokens que minimice la pérdida con un presupuesto fijo.

Suponemos que ambos siguen una relación de ley de potencia: N_opt∝C^α, D_opt∝C^β, y nos gustaría estimar los exponentes desconocidos α y β mediante los siguientes pasos:

Tome el logaritmo de las cantidades: log?(N_opt)=αlog?(C)+const, log?(D_opt)=βlog?(C)+const. Ajustar una regresión lineal. La pendiente de la regresión no es más que el exponente de la ley potencial.

El siguiente código proporciona dicha regresión:

# Ajustar regresión lineal log-log 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))

En nuestro experimento con juguetes, encontramos que N_opt ~C^0.14 y D_opt~ C^0.86. Es posible que este resultado no revele la imagen completa porque hicimos el experimento con modelos y configuraciones simplificados. Pero todavía podemos ver que el crecimiento de la informática conduce a un aumento en el tamaño óptimo del modelo, pero a un ritmo decreciente. Claramente, el presupuesto restante debería atribuirse a más tokens de capacitación.

Además, el cálculo anterior muestra el hecho de que la mejor relación N_opt/D_opt=C^-0,72. Esto implica que cuando aumenta la computación, debe agregar más tokens de entrenamiento en lugar de aumentar el tamaño del modelo.

Conclusiones prácticas

De este experimento, aunque se trata de un caso de juguete, podemos extraer varias ideas:

Para un presupuesto fijo, utilizar un modelo mediano con más datos puede superar a un modelo muy grande con datos limitados. El tamaño óptimo del modelo y el tamaño de los datos crecen con la computación. No entrenes un modelo con muchos parámetros si tienes un presupuesto pequeño. Cuando el presupuesto aumenta, considere primero la relación óptima N_opt/D_opt para determinar si debe aumentar el tamaño del modelo o agregar más datos de entrenamiento.

Conclusión

En esta publicación de blog, proporcionamos un estudio de la compensación entre el tamaño del modelo y los datos bajo un presupuesto de computación fijo para LLM con un estuche de juguete. El experimento muestra que podemos encontrar el par óptimo de tamaño de modelo y número de tokens para lograr el mejor rendimiento del modelo con un presupuesto determinado, lo que permite a los investigadores y profesionales diseñar LLM de manera inteligente y lograr los mejores resultados.

Referencia

[1] Kaplan, J., McCandlish, S., Henighan, T., Brown, TB, Chess, B., Child, R., Gray, S., Radford, A., Wu, J. y Amodei, D. (2020). Leyes de escala para modelos de lenguaje 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). Entrenamiento de modelos de lenguaje grande óptimos para la computación.