Dans l'esprit du programme pour , cet article (et peut-être d'autres à suivre dans une série) étudie les « algorithmes » appris par un modèle de transformateur pour aborder une tâche technique étroite – une version modifiée des « Parenthèses valides » Problème de Leetcode. d'interprétabilité mécaniste les réseaux de neurones Bien que l'utilité de la tâche soit de portée beaucoup plus modeste que la prédiction plus générale du prochain jeton que l'on attend dans un LLM, l'exercice nous aidera à explorer certaines des premières intuitions, outils d'investigation et méthodologies épistémologiques générales généralement déployées pour savoir ce que font les modèles (et comment nous savons qu'ils le font.) Les défis mensuels ARENA Mechinterp ont eu une énorme influence sur ce poste, et la première série de problèmes viendra de là. (Vous devriez absolument consulter .) le programme Structure de la série : Choisissez un problème Leetcode comme tâche. (Partie 1) Entraînez-y un modèle de transformateur minimum viable. (Partie 2) Examinez ce que le modèle a appris. (Partie 3) Partie 1 : Problème Le problème des parenthèses valides vu sur Leetcode : Quelques contraintes modifiées sur le problème que nous utiliserons pour la tâche : Les seuls caractères acceptables sont « ( » et « ) » Cela supprime le besoin de gérer des cas tels que "([)]". La séquence de saisie maximale est de 40 caractères. Pour aider à garder notre modèle petit pour des itérations rapides. Exemples «(((())))» → Valide «()()()(» → Invalide «)()()()(» → Invalide Solution vanille def isValid(self, s: str) -> bool: nesting_depth = 0 for bracket in s: if bracket == '(': # An opening bracket increases unresolved nesting depth nesting_depth += 1 elif bracket == ')': # A closing bracket decreases unresolved nesting depth nesting_depth -= 1 # We don't expect to ever have negative unresolved nesting depth, # so we can declare 'invalid' midway through the sequence if we see this if nesting_depth < 0: return False # Final check that all open brackets were closed. return nesting_depth == 0 Notes sur les cas d'échec : nesting_degree ≠ 0 à la fin de la séquence «()()()((» → Invalide Pour cela, il n'est pas évident que quelque chose ne va pas jusqu'à la toute fin lorsque nous voyons que les dernières parenthèses ouvertes n'ont pas de parenthèse qui les accompagne. La chose à noter est qu'il n'y a aucun moment dans la séquence, jusqu'à la toute fin, où nous avions suffisamment d'informations pour savoir que quelque chose n'allait pas. nesting_degree < 0 à tout moment de la séquence exemple : « ())()()( » → Invalide Dans ce cas, en revanche, il y a suffisamment d'informations en troisième position pour savoir que la validité de la séquence est irrécupérable, nous pouvons donc l'arrêter plus tôt. Il convient de noter que cet exemple aurait réussi le premier test d'échec, car la à la fin aurait été égale à 0. Ce scénario de test ne nous aide donc pas seulement à nous arrêter plus tôt, il est vital. Il en va de même pour le premier exemple de cas d’échec où il aurait réussi le test 2. nesting_depth Maintenant, nous ne nous attendons pas à ce qu'un modèle de transformateur autorégressif résolve le problème exactement de la même manière, étant donné que son architecture offre des mécanismes légèrement différents de ceux qui consistent à parcourir la séquence une fois et à vérifier si tout va bien. Cependant, nous savons avec certitude que l'architecture du transformateur (et d'autres architectures de traitement de séquence) est au moins capable de et des informations sur tous les éléments d'une séquence. Il est important de se rappeler que même si la solution peut paraître différente, la structure du problème est la même et les limites strictes de ce qui est connu et de l'endroit dans la séquence continuent d'être vraies, qu'il s'agisse d'une boucle et d'instructions if ou d'un ensemble d'instructions self. -les balayages d'attention et les non-linéarités MLP. découvrir de traiter La question intéressante est alors de savoir comment cette architecture exploite ces informations et si elles sont facilement perceptibles avec les outils existants ; car il est qu'une solution suffisamment performante de n'importe quelle architecture ne teste pas les deux cas d'échec ci-dessus. inévitable au moins C'est l'un des avantages des problèmes de jouets ; nous obtenons une tâche étroite et suffisamment comprise avec ces garanties fermes qui peuvent contribuer à éclairer l'enquête, comme nous le verrons bientôt. Partie 2 : Données et modèle Préparation des données de formation Voici quelques caractéristiques cibles que nous recherchons avec la génération de données : Un nombre égal de cordes équilibrées et déséquilibrées. Les chaînes seront de longueur paire, car une chaîne de longueur impaire est évidemment déséquilibrée ; ce qui ne serait pas une heuristique très intéressante à apprendre pour le modèle. Toutes les longueurs de chaîne (2 à 40) doivent être également probables. Pour une longueur de chaîne donnée, toutes les profondeurs d'imbrication potentielles des parenthèses doivent être également probables. Un thème commun apparaît : nous essayons de rendre chaque statistique de distribution imaginable également susceptible de réduire les biais dans une direction donnée, de garantir la robustesse et de refuser les heuristiques évidentes à gain rapide comme option pour le modèle. Pour générer des cas d'échec, nous allons d'abord générer des parenthèses valides avec les garanties énumérées ci-dessus, puis en modifier la moitié pour les déséquilibrer. from random import randint, randrange, sample from typing import List, Tuple, Union, Optional, Callable, Dict from jaxtyping import Float, Int import torch as t from torch import Tensor import plotly.express as px import einops from dataclasses import dataclass import math def isValid(s: str) -> bool: nesting_depth = 0 for bracket in s: if bracket == '(': # An opening bracket increases unresolved nesting depth nesting_depth += 1 elif bracket == ')': # A closing bracket decreases unresolved nesting depth nesting_depth -= 1 # We don't expect to ever have negative unresolved nesting depth, # so we can declare 'invalid' midway through the sequence if we see this if nesting_depth < 0: return False # Final check that all open brackets were closed. return nesting_depth == 0 assert isValid('()()((((()())())))') == True assert isValid(')()((((()())()))(') == False Schéma de génération de données n°1 : marche aléatoire La première tentative de génération de parenthèses effectue simplement une marche aléatoire. Mais comme vous pouvez le voir dans les graphiques ci-dessous, le sous-espace des parenthèses déséquilibrées est beaucoup plus grand que celui des parenthèses équilibrées ; nous devrons donc introduire la stochasticité différemment. PARENS = ['(', ')'] def get_random_walk_parens(parens_num: int, length_range: Tuple[int]) -> List[str]: range_start, range_end = length_range random_parens = [ # Add 1 to make passed range_end inclusive ''.join(PARENS[randint(0, 1)] for _ in range(randrange(range_start, range_end + 1, 2))) for _ in range(parens_num) ] return random_parens random_parens = get_random_walk_parens(1000, (2, 10)) random_parens[:10] # output [')(', '(((())()', ')(((()()))', '))))))', '))())()(', '))', '(())', ')()(()()()', ')()())))((', '()'] is_valid_evals = [str(isValid(random_paren)) for random_paren in random_parens] len_evals = [len(random_paren) for random_paren in random_parens] fig = px.histogram(is_valid_evals, title="Count of is-balanced for random walk parentheses strings") fig.show() Schéma de génération de données n°2 : séquence d'imbrication aléatoire gourmande Nous pouvons décomposer la construction d’une chaîne de parenthèses équilibrées en unités discrètes de parenthèses imbriquées. Pour cette construction gourmande, à chaque étape du processus de génération d'une corde une profondeur d'imbrication est choisie parmi un panier de profondeurs viables (pour respecter la longueur de corde cible.) Par exemple, pour la longueur cible , les décompositions d'imbrication uniques suivantes sont possibles : 6 -> [2, 1], [1, 2], [1,1,1] or [3] Corresponding to: -> (())(), ()(()), ()()(), ((())) def get_balanced_parens(nest_depth: int) -> str: """Generate parentheses at the required nesting depth.""" return (PARENS[0] * nest_depth) + (PARENS[1] * nest_depth) assert get_balanced_parens(3) == '((()))' def get_balanced_sequence_parens(nest_depth_sequence: List[int]) -> str: """Return a parentheses string following the nesting depth sequence from a given list.""" return ''.join(get_balanced_parens(nest_depth) for nest_depth in nest_depth_sequence) assert get_balanced_sequence_parens([1,1,2,3]) == '()()(())((()))' def get_random_depth_sequence(target_paren_len: int) -> List[int]: depth_sequence = [] while target_paren_len > 0: depth = randint(1, target_paren_len / 2) depth_sequence.append(depth) target_paren_len -= 2 * depth return depth_sequence rand_depth_seq = get_random_depth_sequence(10) print(rand_depth_seq) # Example output: '[3, 1, 1]' assert sum([2 * depth for depth in rand_depth_seq]) == 10 def get_random_sequence_parens(parens_num: int, length_range: Tuple[int]) -> List[str]: random_depth_sequences = [get_random_depth_sequence( randrange(*length_range, 2) ) for _ in range(parens_num)] random_parens = [ get_balanced_sequence_parens(random_depth_sequence) for random_depth_sequence in random_depth_sequences ] return random_parens, random_depth_sequences Obtenez des parenthèses équilibrées random_seq_parens, depth_sequences = get_random_sequence_parens(100000, (2, 11)) is_valid_evals = [str(isValid(random_paren)) for random_paren in random_seq_parens] len_evals = [len(random_paren) for random_paren in random_seq_parens] Voyons les fréquences des profondeurs de nidification depth_freq = {} for seq in depth_sequences: for depth in seq: depth_freq.setdefault(depth, 0) depth_freq[depth] += 1 depth_freq # output -> {2: 39814, 1: 100088, 3: 20127, 4: 9908, 5: 4012} depth_seq_hist = px.histogram(depth_sequences, title="Frequence of nesting depths in 'Random Nesting Depth Sequence' Output") depth_seq_hist.show() Et maintenant, voyons les fréquences de longueur. paren_len_hist = px.histogram(len_evals, title="Frequency of string lengths") paren_len_hist.show() Remarque sur la génération de données Notez qu'il existe une tension entre les propriétés potentielles suivantes de notre distribution de données. Chaque longueur de chaîne est également probable. Chaque sous-chaîne de profondeur d'imbrication a la même probabilité sur toutes les chaînes. En effet, les sous-séquences d'imbrication à faible profondeur auront plus de chances d'apparaître dans une séquence d'imbrication aléatoire donnée, comme le montrent les graphiques ci-dessus. Pour contrer cette tendance naturelle de la séquence purement aléatoire, lors de la génération d'une sous-chaîne donnée de parenthèses, nous pourrions échantillonner à partir d'une distribution asymétrique pour rendre plus probables les valeurs d'imbrication plus profondes. Celui-ci sera revisité après un premier passage à l'entraînement. px.histogram(random_seq_parens, title="Frequency of balanced Parentheses").show() Créer des parenthèses déséquilibrées Notre ensemble de données ne peut pas avoir uniquement des parenthèses équilibrées. Nous pouvons donc créer une stratégie de génération de données pour dériver des chaînes déséquilibrées de notre ensemble de données équilibré. def _flip_idx(idx): return (idx + 1) % 2 assert _flip_idx(0) == 1 assert _flip_idx(1) == 0 def make_parens_unbalanced(paren: str) -> str: """Take balanced-parentheses and randomly mutate it till it's unbalanced. Both the number of mutations and indices are chosen at random. """ paren_idx_dict = {'(': 0, ')': 1} paren_list = list(paren) num_flipped_positions = randint(1, len(paren)) while isValid(''.join(paren_list)): flip_points = sample(range(len(paren)), num_flipped_positions) for flip_idx in flip_points: idx_char = paren_idx_dict[paren_list[flip_idx]] flipped_idx = _flip_idx(idx_char) paren_list[flip_idx] = PARENS[flipped_idx] return ''.join(paren_list) assert not isValid(make_parens_unbalanced('((()))')) Obtenir un ensemble de données Parens déséquilibré unbal_random_seq_parens = [make_parens_unbalanced(paren) for paren in random_seq_parens] Formation sur modèle Maintenant que nous avons nos ensembles de données, pour le plaisir, nous allons écrire notre architecture Transformer à partir de zéro. D'abord quelques configurations @dataclass class Config: context_len = 12 d_vocab: int = 5 d_out_vocab: int = 2 d_model: int = 56 d_head = 28 d_mlp = 56 * 4 causal_attention = False num_heads = 2 num_layers = 3 init_range: float = 1 PAD_TOKEN_IDX = 1 Ensuite, notre tokenizer pour analyser les entrées : class Tokenizer: def __init__(self, vocab: str, context_width: Int, enforce_context: bool=False): self.START_TOKEN, START_TOKEN_IDX = "<start>", 0 self.PAD_TOKEN, PAD_TOKEN_IDX = "<pad>", 1 self.END_TOKEN, END_TOKEN_IDX = "<end>", 2 util_tokens_t_to_i = {self.START_TOKEN: START_TOKEN_IDX, self.PAD_TOKEN: PAD_TOKEN_IDX, self.END_TOKEN: END_TOKEN_IDX} util_tokens_i_to_t = {START_TOKEN_IDX: self.START_TOKEN, PAD_TOKEN_IDX: self.PAD_TOKEN, END_TOKEN_IDX: self.END_TOKEN} self.enforce_context = enforce_context self.context_width = context_width self.vocab = vocab self.t_to_i = {**util_tokens_t_to_i, **{token: token_id + 3 for token_id, token in enumerate(self.vocab)}} self.i_to_t = {**util_tokens_i_to_t, **{token_id + 3: token for token_id, token in enumerate(self.vocab)}} @staticmethod def pad_sequence(sequence: str, end_token: str, pad_token: str, max_length: Int, enforce_context: bool) -> List[str]: if not enforce_context: # Truncate if sequence length is greater sequence = sequence[:max_length] else: assert len(sequence) <= max_length, f"Sequence length is greater than the max allowed data length: {max_length}" return list(sequence) + [end_token] + [pad_token] * (max_length - len(sequence)) def tokenize(self, data: Union[str, List[str]]) -> Int[Tensor, "batch seq"]: if isinstance(data, str): data = [data] def _list_tokens_to_id(tokens: List[str]) -> List[Int]: return [self.t_to_i[token] for token in tokens] # to leave room for start and end tokens max_seq_len = self.context_width - 2 data_as_tokens = [ _list_tokens_to_id([ self.START_TOKEN, *self.pad_sequence(seq, self.END_TOKEN, self.PAD_TOKEN, max_seq_len, self.enforce_context), ]) for seq in data ] return t.tensor(data_as_tokens) (Dé)Incrustations class EmbedLayer(t.nn.Module): def __init__(self, cfg: Config): super().__init__() self.W_E = t.nn.Parameter(t.empty(cfg.d_vocab, cfg.d_model)) t.nn.init.normal_(self.W_E, mean=0.0, std=cfg.init_range) def forward(self, x: Int[Tensor, "batch seq"]) -> Int[Tensor, "batch seq d_model"]: return self.W_E[x] class UnEmbedLayer(t.nn.Module): def __init__(self, cfg: Config): super().__init__() self.W_U = t.nn.Parameter(t.empty(cfg.d_model, cfg.d_out_vocab)) t.nn.init.normal_(self.W_U, mean=0.0, std=cfg.init_range) def forward(self, x: Int[Tensor, "batch seq d_model"]) -> Int[Tensor, "batch seq d_out_vocab"]: return x @ self.W_U class PositionalEmbedding(t.nn.Module): def __init__(self, cfg: Config): super().__init__() denom = t.exp( t.arange(0, cfg.d_model, 2) * -(math.log(10000.0) / cfg.d_model) ) pos = t.arange(0, cfg.context_len).unsqueeze(1) param = pos * denom P_E = t.zeros(cfg.context_len, cfg.d_model) P_E[:, 0::2] = t.sin(param) P_E[:, 1::2] = t.cos(param) P_E = P_E.unsqueeze(0) self.register_buffer("P_E", P_E) def forward(self, x): _batch, seq_len, d_model = x.shape x = x + self.P_E[..., :seq_len, :d_model].requires_grad_(False) return x Norme de couche pratique class LayerNorm(t.nn.Module): def __init__(self, cfg): super().__init__() self.scale = t.nn.Parameter(t.ones(cfg.d_model)) self.bias = t.nn.Parameter(t.zeros(cfg.d_model)) def forward(self, x): mean = t.mean(x, dim=2, keepdim=True) var = t.var(x, dim=2, keepdim=True, unbiased=False) y = (x - mean) / (var + 0.00001).sqrt() return (y * self.scale) + self.bias Et enfin Attention ! class AttentionLayer(t.nn.Module): def __init__(self, cfg): super().__init__() self.register_buffer("IGNORE", t.tensor(-1e5, dtype=t.float32)) self.cfg = cfg self.W_Q = t.nn.Parameter(t.empty(cfg.num_heads, cfg.d_model, cfg.d_head)) self.W_K = t.nn.Parameter(t.empty(cfg.num_heads, cfg.d_model, cfg.d_head)) self.W_V = t.nn.Parameter(t.empty(cfg.num_heads, cfg.d_model, cfg.d_head)) self.W_O = t.nn.Parameter(t.empty(cfg.num_heads, cfg.d_head, cfg.d_model)) self.b_Q = t.nn.Parameter(t.zeros(cfg.num_heads, cfg.d_head)) self.b_K = t.nn.Parameter(t.zeros(cfg.num_heads, cfg.d_head)) self.b_V = t.nn.Parameter(t.zeros(cfg.num_heads, cfg.d_head)) self.b_O = t.nn.Parameter(t.zeros(cfg.d_model)) t.nn.init.normal_(self.W_Q, mean=0.0, std=cfg.init_range) t.nn.init.normal_(self.W_K, mean=0.0, std=cfg.init_range) t.nn.init.normal_(self.W_V, mean=0.0, std=cfg.init_range) t.nn.init.normal_(self.W_O, mean=0.0, std=cfg.init_range) def forward(self, params): #TODO: revisit implementing pad_mask with hooks x, pad_mask = params Q = einops.einsum(x, self.W_Q, 'bs dm, h dm dh -> bsh dh') + self.b_Q K = einops.einsum(x, self.W_K, 'bs dm, h dm dh -> bsh dh') + self.b_K V = einops.einsum(x, self.W_V, 'bs dm, h dm dh -> bsh dh') + self.b_V attention_scores = einops.einsum(Q, K, 'b s_q h dh, b s_k h dh -> bh s_q s_k') scaled_attention_scores = attention_scores / (self.cfg.d_head ** 0.5) if self.cfg.causal_attention: scaled_attention_scores = self.apply_causal_mask(scaled_attention_scores) scaled_attention_scores = self.apply_padding_mask(scaled_attention_scores, pad_mask) attention_patterns = t.nn.Softmax(dim=-1)(scaled_attention_scores) post_attention_values = einops.einsum( attention_patterns, V, 'bh s_q s_k, b s_k h dh -> b s_q h dh' ) out = einops.einsum( post_attention_values, self.W_O, 'b s_q h dh, h dh dm -> b s_q dm' ) + self.b_O return out def apply_causal_mask(self, attention_scores): b, h, s_q, s_k = attention_scores.shape mask = t.tril(t.ones(s_q,s_k)).bool() return t.where(mask, attention_scores, self.IGNORE) def apply_padding_mask(self, attention_scores, pad_mask): return t.where(pad_mask, attention_scores, self.IGNORE) Couches MLP class LinearLayer(t.nn.Module): def __init__(self, in_dim, out_dim, include_bias=True): super().__init__() self.include_bias = include_bias self.W = t.nn.Parameter(t.empty(in_dim, out_dim)) t.nn.init.normal_(self.W, mean=0.0, std=cfg.init_range) self.b = None if include_bias: self.b = t.zeros(out_dim) def forward(self, x: Int[Tensor, "batch seq in_dim"]) -> Int[Tensor, "batch seq out_dim"]: out = x @ self.W if self.include_bias: out = out + self.b return out class MLP(t.nn.Module): def __init__(self, cfg): super().__init__() self.in_layer = LinearLayer(cfg.d_model, cfg.d_mlp) self.out_layer = LinearLayer(cfg.d_mlp, cfg.d_model) self.non_linearity = t.nn.ReLU() def forward(self, x): post_W_in = self.in_layer(x) post_non_lin = self.non_linearity(post_W_in) return self.out_layer(post_non_lin) L'assembler dans un transformateur class TransformerBlock(t.nn.Module): def __init__(self, cfg): super().__init__() self.ln1 = LayerNorm(cfg) self.attention = AttentionLayer(cfg) self.ln2 = LayerNorm(cfg) self.mlp = MLP(cfg) def forward(self, params): x, pad_mask = params resid_mid = self.attention((self.ln1(x), pad_mask)) + x resid_post = self.mlp(self.ln2(resid_mid)) + resid_mid return resid_post, pad_mask class Transformer(t.nn.Module): def __init__(self, cfg: Config): super().__init__() self.cfg = cfg self.embed = EmbedLayer(cfg) self.pos_embed = PositionalEmbedding(cfg) self.final_ln = LayerNorm(cfg) self.unembed = UnEmbedLayer(cfg) self.blocks = t.nn.Sequential(*([TransformerBlock(cfg)] * cfg.num_layers)) def forward(self, x): #TODO: revisit implementing pad_mask with hooks pad_mask = self.get_pad_mask(x) res_post_pos_embed = self.pos_embed(self.embed(x)) post_blocks, _ = self.blocks((res_post_pos_embed, pad_mask)) logits = self.unembed(self.final_ln(post_blocks)) return logits def get_pad_mask(self, x): batch, seq = x.shape return einops.repeat(x != self.cfg.PAD_TOKEN_IDX, 'batch seq -> batch 1 seq_q seq', seq_q=seq) Utilitaires de formation def cross_entropy_loss(output, targets): log_probs = output.log_softmax(dim=-1) predictions = log_probs[:, 0] batch, out_dim = predictions.shape true_output = predictions[range(batch), targets] return -true_output.sum() / batch def test(model, data, loss_func): inputs, targets = data with t.no_grad(): output = model(inputs) loss = loss_func(output, targets) return loss def train(model, data, optimizer, loss_func): inputs, targets = data optimizer.zero_grad() output = model(inputs) loss = loss_func(output, targets) loss.backward() optimizer.step() return loss Configuration de la formation cfg = Config() tokenizer = Tokenizer('()', 12, True) inputs = tokenizer.tokenize([*unbal_random_seq_parens, *random_seq_parens]) targets = t.tensor([*([0] * len(unbal_random_seq_parens)), *([1] * len(random_seq_parens))]) rand_indices = t.randperm(targets.shape[0]) rand_inputs = inputs[rand_indices, :] rand_targets = targets[rand_indices] model = Transformer(cfg) adamW = t.optim.AdamW(model.parameters(), lr=0.01) Formation réelle batch_size = 10000 train_size = int(0.7 * batch_size) epochs = 15 for epoch in range(epochs): for batch_id in range(0, rand_inputs.shape[0], batch_size): rand_inputs_batch, rand_targets_batch = rand_inputs[batch_id : batch_id + batch_size], rand_targets[batch_id : batch_id + batch_size] train_input, train_target = rand_inputs_batch[:train_size, :], rand_targets_batch[:train_size] test_input, test_target = rand_inputs_batch[train_size:, :], rand_targets_batch[train_size:] train(model, (train_input, train_target), adamW, cross_entropy_loss) test_loss = test(model, (test_input, test_target), cross_entropy_loss) print(f'Loss: {test_loss} on epoch: {epoch}/{epochs}') Dans la troisième partie, nous étudierons les composants internes de ce réseau formé. Nous le ferons en examinant les modèles d'attention et en appliquant certains des outils de diagnostic de l'interprétabilité mécaniste, tels que les correctifs d'activation, pour construire un modèle mécaniste permettant de comprendre comment le réseau a résolu cette tâche. Merci d'avoir lu jusqu'ici et à bientôt dans la troisième partie !
