El desarrollo de modelos de lenguaje a gran escala, especialmente ChatGPT, ha dejado a quienes han experimentado con él, incluido yo mismo, asombrados por su notable destreza lingüística y su capacidad para realizar diversas tareas. Sin embargo, muchos investigadores, incluido yo mismo, aunque se maravillan de sus capacidades, también se quedan perplejos. A pesar de conocer la arquitectura del modelo y los valores específicos de sus pesos, todavía nos cuesta comprender por qué una secuencia particular de entradas conduce a una secuencia específica de salidas.
En esta publicación de blog, intentaré desmitificar GPT2-small utilizando la interpretabilidad mecanicista en un caso simple: la predicción de tokens repetidos.
Las herramientas matemáticas tradicionales para explicar los modelos de aprendizaje automático no son del todo adecuadas para los modelos de lenguaje.
Considere SHAP, una herramienta útil para explicar los modelos de aprendizaje automático. Es competente para determinar qué característica influyó significativamente en la predicción de un vino de buena calidad. Sin embargo, es importante recordar que los modelos de lenguaje hacen predicciones a nivel de token, mientras que los valores SHAP se calculan principalmente a nivel de característica, lo que los hace potencialmente no aptos para tokens.
Además, los modelos de lenguaje (LLM) tienen numerosos parámetros y entradas, lo que crea un espacio de alta dimensión. Calcular los valores SHAP es costoso incluso en espacios de baja dimensión, y más aún en el espacio de alta dimensión de los LLM.
A pesar de tolerar los altos costos computacionales, las explicaciones proporcionadas por SHAP pueden ser superficiales. Por ejemplo, saber que el término “potter” influyó más en la predicción de resultados debido a la mención anterior de “Harry” no proporciona mucha información. Nos deja inseguros sobre la parte del modelo o el mecanismo específico responsable de tal predicción.
La interpretabilidad mecanicista ofrece un enfoque diferente. No solo identifica características o entradas importantes para las predicciones de un modelo. Más bien, arroja luz sobre los mecanismos subyacentes o procesos de razonamiento, ayudándonos a comprender cómo un modelo hace sus predicciones o decisiones.
Usaremos GPT2-small para una tarea sencilla: predecir una secuencia de tokens repetidos. La biblioteca que usaremos es Lente transformadoraque está diseñado para interpretabilidad mecanicista de modelos de lenguaje de estilo GPT-2.
gpt2_small: HookedTransformer = HookedTransformer.from_pretrained("gpt2-small")
Usamos el código anterior para cargar el modelo GPT2-Small y predecir tokens en una secuencia generada por una función específica. Esta secuencia incluye dos secuencias de tokens idénticas, seguidas por bos_token. Un ejemplo sería “ABCDABCD” + bos_token cuando seq_len es 3. Para mayor claridad, nos referimos a la secuencia desde el principio hasta seq_len como la primera mitad, y a la secuencia restante, excluyendo bos_token, como la segunda mitad.
def generate_repeated_tokens(
model: HookedTransformer, seq_len: int, batch: int = 1
) -> Int[Tensor, "batch full_seq_len"]:
'''
Generates a sequence of repeated random tokensOutputs are:
rep_tokens: [batch, 1+2*seq_len]
'''
bos_token = (t.ones(batch, 1) * model.tokenizer.bos_token_id).long() # generate bos token for each batch
rep_tokens_half = t.randint(0, model.cfg.d_vocab, (batch, seq_len), dtype=t.int64)
rep_tokens = t.cat([bos_token,rep_tokens_half,rep_tokens_half], dim=-1).to(device)
return rep_tokens
Cuando permitimos que el modelo se ejecute en el token generado, encontramos una observación interesante: el modelo funciona significativamente mejor en la segunda mitad de la secuencia que en la primera mitad. Esto se mide mediante las probabilidades logarítmicas de los tokens correctos. Para ser precisos, el rendimiento en la primera mitad es -13,898, mientras que el rendimiento en la segunda mitad es -0,644.
También podemos calcular la precisión de la predicción, definida como la proporción de tokens predichos correctamente (aquellos idénticos a los tokens generados) con respecto al número total de tokens. La precisión de la primera mitad de la secuencia es 0,0, lo cual no es sorprendente ya que estamos trabajando con tokens aleatorios que carecen de significado real. Mientras tanto, la precisión para la segunda mitad es de 0,93, superando significativamente a la primera mitad.
Encontrar cabeza de inducción
La observación anterior podría explicarse por la existencia de un circuito de inducción. Este es un circuito que escanea la secuencia en busca de instancias anteriores del token actual, identifica el token que lo siguió anteriormente y predice que la misma secuencia se repetirá. Por ejemplo, si encuentra una ‘A’, busca la ‘A’ anterior o un token muy similar a ‘A’ en el espacio de incrustación, identifica el token posterior ‘B’ y luego predice el siguiente token después de ‘A’. ‘ para ser ‘B’ o un token muy similar a ‘B’ en el espacio de incrustación.
Este proceso de predicción se puede dividir en dos pasos:
- Identifique el token anterior igual (o similar). Cada token en la segunda mitad de la secuencia debe “prestar atención” al token ‘seq_len’ que se encuentra antes. Por ejemplo, la ‘A’ en la posición 4 debe prestar atención a la ‘A’ en la posición 1 si ‘seq_len’ es 3. Podemos llamar la atención a la cabeza que realiza esta tarea el “cabeza de inducción.”
- Identifique el siguiente token ‘B’. Este es el proceso de copiar información del token anterior (por ejemplo, ‘A’) al siguiente token (por ejemplo, ‘B’). Esta información se utilizará para “reproducir” ‘B’ cuando ‘A’ aparezca nuevamente. Podemos llamar la atención del jefe que realiza esta tarea el “cabeza de token anterior.”
Estos dos cabezales constituyen un circuito de inducción completo. Tenga en cuenta que a veces el término “cabezal de inducción” también se utiliza para describir todo el “circuito de inducción”. Para obtener más información sobre el circuito de inducción, recomiendo encarecidamente el artículo. Jefe de inducción y aprendizaje en contexto que es una obra maestra!
Ahora, identifiquemos el cabezal de atención y el cabezal anterior en GPT2-small.
El siguiente código se utiliza para encontrar el cabezal de inducción. Primero, ejecutamos el modelo con 30 lotes. Luego, calculamos el valor medio de la diagonal con un desplazamiento de seq_len en la matriz del patrón de atención. Este método nos permite medir el grado de atención que le da el token actual al que aparece seq_len de antemano.
def induction_score_hook(
pattern: Float[Tensor, "batch head_index dest_pos source_pos"],
hook: HookPoint,
):
'''
Calculates the induction score, and stores it in the [layer, head] position of the `induction_score_store` tensor.
'''
induction_stripe = pattern.diagonal(dim1=-2, dim2=-1, offset=1-seq_len) # src_pos, des_pos, one position right from seq_len
induction_score = einops.reduce(induction_stripe, "batch head_index position -> head_index", "mean")
induction_score_store[hook.layer(), :] = induction_scoreseq_len = 50
batch = 30
rep_tokens_30 = generate_repeated_tokens(gpt2_small, seq_len, batch)
induction_score_store = t.zeros((gpt2_small.cfg.n_layers, gpt2_small.cfg.n_heads), device=gpt2_small.cfg.device)
rep_tokens_30,
return_type=None,
pattern_hook_names_filter,
induction_score_hook
)]
)
Ahora, examinemos los puntajes de inducción. Notaremos que algunas cabezas, como la de la capa 5 y la cabeza 5, tienen una puntuación de inducción alta de 0,91.
También podemos mostrar el patrón de atención de esta cabeza. Notará una línea diagonal clara hasta un desplazamiento de seq_len.
De manera similar, podemos identificar el encabezado del token anterior. Por ejemplo, la capa 4, cabeza 11, demuestra un patrón fuerte para el token anterior.
¿Cómo se atribuyen las capas MLP?
Consideremos esta pregunta: ¿cuentan las capas MLP? Sabemos que GPT2-Small contiene capas de atención y MLP. Para investigar esto, propongo utilizar una técnica de ablación.
La ablación, como su nombre lo indica, elimina sistemáticamente ciertos componentes del modelo y observa cómo cambia el rendimiento como resultado.
Reemplazaremos la salida de las capas MLP en la segunda mitad de la secuencia con las de la primera mitad y observaremos cómo esto afecta la función de pérdida final. Calcularemos la diferencia entre la pérdida después de reemplazar las salidas de la capa MLP y la pérdida original de la segunda mitad de la secuencia usando el siguiente código.
def patch_residual_component(
residual_component,
hook,
pos,
cache,
):
residual_component[0,pos, :] = cache[hook.name][pos-seq_len, :]
return residual_componentablation_scores = t.zeros((gpt2_small.cfg.n_layers, seq_len), device=gpt2_small.cfg.device)
gpt2_small.reset_hooks()
logits = gpt2_small(rep_tokens, return_type="logits")
loss_no_ablation = cross_entropy_loss(logits[:, seq_len: max_len],rep_tokens[:, seq_len: max_len])
for layer in tqdm(range(gpt2_small.cfg.n_layers)):
for position in range(seq_len, max_len):
hook_fn = functools.partial(patch_residual_component, pos=position, cache=rep_cache)
ablated_logits = gpt2_small.run_with_hooks(rep_tokens, fwd_hooks=[
(utils.get_act_name("mlp_out", layer), hook_fn)
])
loss = cross_entropy_loss(ablated_logits[:, seq_len: max_len], rep_tokens[:, seq_len: max_len])
ablation_scores[layer, position-seq_len] = loss - loss_no_ablation
Llegamos a un resultado sorprendente: aparte del primer token, la ablación no produce una diferencia logit significativa. Esto sugiere que las capas MLP pueden no tener una contribución significativa en el caso de tokens repetidos.
Dado que las capas MLP no contribuyen significativamente a la predicción final, podemos construir manualmente un circuito de inducción usando el cabezal de la capa 5, cabezal 5, y el cabezal de la capa 4, cabezal 11. Recuerde que estos son el cabezal de inducción y la cabeza del token anterior. Lo hacemos mediante el siguiente código:
def K_comp_full_circuit(
model: HookedTransformer,
prev_token_layer_index: int,
ind_layer_index: int,
prev_token_head_index: int,
ind_head_index: int
) -> FactoredMatrix:
'''
Returns a (vocab, vocab)-size FactoredMatrix,
with the first dimension being the query side
and the second dimension being the key side (going via the previous token head)'''
W_E = gpt2_small.W_E
W_Q = gpt2_small.W_Q[ind_layer_index, ind_head_index]
W_K = model.W_K[ind_layer_index, ind_head_index]
W_O = model.W_O[prev_token_layer_index, prev_token_head_index]
W_V = model.W_V[prev_token_layer_index, prev_token_head_index]
Q = W_E @ W_Q
K = W_E @ W_V @ W_O @ W_K
return FactoredMatrix(Q, K.T)
Calcular la precisión superior de este circuito produce un valor de 0,2283. ¡Esto es bastante bueno para un circuito construido con sólo dos cabezas!
Para una implementación detallada, consulte mi computadora portátil.