Interpretabilità meccanicistica sulla previsione di token ripetuti
2 ore fa
Lo sviluppo di modelli linguistici su larga scala, in particolare ChatGPT, ha lasciato coloro che lo hanno sperimentato, me compreso, stupiti dalla sua notevole abilità linguistica e dalla sua capacità di svolgere compiti diversi. Tuttavia, molti ricercatori, me compreso, pur meravigliandosi delle sue capacità, si ritrovano anche perplessi. Pur conoscendo l’architettura del modello e i valori specifici dei suoi pesi, facciamo ancora fatica a comprendere perché una particolare sequenza di input porta a una specifica sequenza di output.
In questo post del blog, cercherò di demistificare GPT2-small utilizzando l'interpretabilità meccanicistica su un caso semplice: la previsione di token ripetuti.
Gli strumenti matematici tradizionali per spiegare i modelli di apprendimento automatico non sono del tutto adatti ai modelli linguistici.
Considera SHAP, uno strumento utile per spiegare i modelli di machine learning. È abile nel determinare quale caratteristica ha influenzato in modo significativo la previsione di un vino di buona qualità. Tuttavia, è importante ricordare che i modelli linguistici effettuano previsioni a livello di token, mentre i valori SHAP sono per lo più calcolati a livello di funzionalità, rendendoli potenzialmente inadatti ai token.
Inoltre, i modelli linguistici (LLM) hanno numerosi parametri e input, creando uno spazio ad alta dimensione. Il calcolo dei valori SHAP è costoso anche negli spazi a bassa dimensionalità, e ancora di più nello spazio ad alta dimensionalità degli LLM.
Nonostante sopportino gli elevati costi computazionali, le spiegazioni fornite da SHAP possono essere superficiali. Ad esempio, sapere che il termine “potter” ha influenzato maggiormente la previsione della produzione a causa della precedente menzione di “Harry” non fornisce molte informazioni. Ciò ci lascia incerti sulla parte del modello o sul meccanismo specifico responsabile di tale previsione.
L’interpretazione meccanicistica offre un approccio diverso. Non si limita a identificare caratteristiche o input importanti per le previsioni di un modello. Invece, fa luce sui meccanismi sottostanti o sui processi di ragionamento, aiutandoci a capire come un modello fa le sue previsioni o decisioni.
Utilizzeremo GPT2-small per un compito semplice: prevedere una sequenza di token ripetuti. La libreria che useremo è Obiettivo trasformatoreper cui è stato progettato interpretabilità meccanicistica di modelli linguistici in stile GPT-2.
gpt2_small: HookedTransformer = HookedTransformer.from_pretrained("gpt2-small")
Usiamo il codice sopra per caricare il modello GPT2-Small e prevedere i token su una sequenza generata da una funzione specifica. Questa sequenza include due sequenze di token identiche, seguite da bos_token. Un esempio potrebbe essere “ABCDABCD” + bos_token quando seq_len è 3. Per chiarezza, ci riferiamo alla sequenza dall'inizio a seq_len come prima metà e alla sequenza rimanente, escluso bos_token, come seconda metà.
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
Quando permettiamo al modello di funzionare sul token generato, troviamo un'osservazione interessante: il modello funziona significativamente meglio nella seconda metà della sequenza rispetto alla prima metà. Questo viene misurato dalle probabilità logaritmiche sui token corretti. Per la precisione la prestazione del primo tempo è di -13.898, mentre quella del secondo tempo è di -0.644.
Possiamo anche calcolare l'accuratezza della previsione, definita come il rapporto tra i token correttamente previsti (quelli identici ai token generati) e il numero totale di token. La precisione per la sequenza della prima metà è 0,0, il che non sorprende poiché stiamo lavorando con token casuali privi di significato effettivo. Nel frattempo, la precisione per il secondo tempo è 0,93, superando significativamente quella del primo tempo.
Trovare la testa di induzione
L'osservazione di cui sopra potrebbe essere spiegata dall'esistenza di un circuito di induzione. Questo è un circuito che analizza la sequenza per individuare le istanze precedenti del token corrente, identifica il token che lo ha seguito in precedenza e prevede che la stessa sequenza si ripeterà. Ad esempio, se incontra una “A”, cerca la “A” precedente o un token molto simile ad “A” nello spazio di incorporamento, identifica il token successivo “B” ', e quindi predice che il token successivo dopo 'A' sarà 'B' o un token molto simile a 'B' nello spazio di incorporamento.
Questo processo di previsione può essere suddiviso in due fasi:
- Identificare il token precedente uguale (o simile). Ogni token nella seconda metà della sequenza dovrebbe “prestare attenzione” al token “seq_len” posto prima di esso. Ad esempio, la “A” in posizione 4 dovrebbe prestare attenzione alla “A” in posizione 1 se “seq_len” è 3. Possiamo chiamare la testa dell'attenzione che esegue questo compito “testa ad induzione.â€
- Identifica il seguente token “B”. Questo è il processo di copia delle informazioni dal token precedente (ad esempio, “A”) nel token successivo (ad esempio, “B”). Queste informazioni verranno utilizzate per “riprodurre” “B” quando “A” appare di nuovo. Possiamo chiamare la testa dell'attenzione che svolge questo compito “testa del token precedente.â€
Queste due teste costituiscono un circuito di induzione completo. Si noti che a volte il termine “testa di induzione” viene utilizzato anche per descrivere l'intero “circuito di induzione”. Per una maggiore introduzione al circuito di induzione, consiglio vivamente l'articolo Responsabile dell'apprendimento e dell'induzione nel contesto che è un capolavoro!
Ora identifichiamo la testa dell'attenzione e la testa precedente in GPT2-small.
Il seguente codice viene utilizzato per trovare la testa di induzione. Innanzitutto, eseguiamo il modello con 30 batch. Quindi, calcoliamo il valore medio della diagonale con un offset di seq_len nella matrice del modello di attenzione. Questo metodo ci permette di misurare il grado di attenzione che il token corrente presta a quello che appare in anticipo seq_len.
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
))
)
Ora esaminiamo i punteggi di induzione. Noteremo che alcune teste, come quella sul livello 5 e la testa 5, hanno un punteggio di induzione elevato pari a 0,91.
Possiamo anche visualizzare il modello di attenzione di questa testa. Noterai una linea diagonale chiara fino a un offset di seq_len.
Allo stesso modo, possiamo identificare la testa del token precedente. Ad esempio, il livello 4, testa 11, mostra un modello forte per il token precedente.
Come si attribuiscono i livelli MLP?
Consideriamo questa domanda: contano i livelli MLP? Sappiamo che GPT2-Small contiene sia livelli di attenzione che MLP. Per indagare su questo, propongo di utilizzare una tecnica di ablazione.
L'ablazione, come suggerisce il nome, rimuove sistematicamente alcuni componenti del modello e osserva come cambiano le prestazioni di conseguenza.
Sostituiremo l'output degli strati MLP nella seconda metà della sequenza con quelli della prima metà e osserveremo come ciò influisce sulla funzione di perdita finale. Calcoleremo la differenza tra la perdita dopo la sostituzione degli output dello strato MLP e la perdita originale della seconda metà della sequenza utilizzando il seguente codice.
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
Arriviamo ad un risultato sorprendente: a parte il primo token, l’ablazione non produce una differenza logit significativa. Ciò suggerisce che gli strati MLP potrebbero non avere un contributo significativo nel caso di token ripetuti.
Dato che gli strati MLP non contribuiscono in modo significativo alla previsione finale, possiamo costruire manualmente un circuito di induzione utilizzando la testa dello strato 5, testa 5, e la testa dello strato 4, testa 11. Ricordiamo che queste sono le teste di induzione e la testa del token precedente. Lo facciamo con il seguente codice:
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)
Calcolando la prima precisione di questo circuito si ottiene un valore di 0,2283. Questo è abbastanza buono per un circuito costruito da sole due teste!
Per un'implementazione dettagliata, consulta il mio taccuino.
Fonte: towardsdatascience.com
