En el espíritu de la agenda para , esta publicación (y quizás otras que seguirán en una serie) investiga los "algoritmos" aprendidos por un modelo transformador para abordar una tarea técnica limitada: una versión modificada de los "Paréntesis válidos". Problema de código Leet. de interpretabilidad mecanicista redes neuronales Si bien la utilidad de la tarea tiene un alcance mucho más modesto que la predicción más general del siguiente token que se esperaría en un LLM, el ejercicio nos ayudará a explorar algunas de las primeras intuiciones, herramientas de investigación y metodologías epistemológicas generales que normalmente se implementan para saber qué están haciendo los modelos (y cómo sabemos que están haciendo). Los desafíos mensuales de ARENA Mechinterp fueron una gran influencia en esta publicación, y el primer conjunto de problemas surgirá de ahí. (Definitivamente deberías consultar ). el programa Estructura de la serie: Elija un problema de Leetcode como tarea. (Parte 1) Entrene un modelo de transformador mínimo viable en él. (Parte 2) Investiga lo que aprendió el modelo. (Parte 3) Parte 1: Problema El problema de los paréntesis válidos como se ve en Leetcode: Algunas restricciones modificadas sobre el problema que usaremos para la tarea: Los únicos caracteres aceptables son "(" y ")" Esto elimina la necesidad de manejar casos como "([)]”. La secuencia de entrada máxima es de 40 caracteres. Para ayudar a mantener nuestro modelo pequeño para iteraciones rápidas. Ejemplos “(((())))” → Válido “()()()(” → Inválido “)()()()(” → Inválido Solución de vainilla 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 Notas sobre casos de falla: nesting_ Depth ≠ 0 al final de la secuencia “()()()((” → No válido Por esto, no es obvio que algo esté mal hasta el final, cuando vemos que los últimos corchetes abiertos no tienen un corchete que los acompañe. Lo que hay que tener en cuenta es que no hay ningún punto en la secuencia, hasta el final, en el que tuviéramos suficiente información para saber que algo estaba mal. nesting_ Depth < 0 en cualquier punto de la secuencia ejemplo: “())()()(” → No válido En este caso, por otro lado, hay suficiente información en la tercera posición para saber que la validez de la secuencia es irrecuperable, por lo que podemos abandonarla antes de tiempo. Algo a tener en cuenta es que este ejemplo habría pasado la primera prueba de falla ya que al final habría sido igual a 0. Entonces, este caso de prueba no solo nos ayuda a detenernos temprano, sino que es vital. Lo mismo se aplica al primer ejemplo de caso de falla en el que habría pasado la prueba 2. nesting_depth Ahora bien, no esperamos que un modelo de transformador autorregresivo resuelva el problema exactamente de la misma manera, dado que su arquitectura ofrece mecanismos ligeramente diferentes a los de recorrer la secuencia una vez y verificar si todo está bien. Sin embargo, sabemos con certeza que la arquitectura del transformador (y otras arquitecturas de procesamiento de secuencias) al menos son capaces de y información sobre todos los elementos de una secuencia. Es importante recordar que, si bien la solución puede parecer diferente, la estructura del problema es la misma y los límites estrictos de lo que se conoce y en qué parte de la secuencia siguen siendo ciertos, ya sea un bucle y declaraciones if o un conjunto de self. -barridos de atención y no linealidades MLP. descubrir procesar La pregunta interesante entonces es cómo aprovecha esta arquitectura esta información y si es fácilmente discernible con las herramientas existentes; porque es que una solución de cualquier arquitectura con suficiente rendimiento no pruebe los dos casos de falla anteriores. inevitable al menos Ésta es una de las ventajas de los problemas con los juguetes; Con estas garantías estrictas obtenemos una tarea suficientemente comprendida y limitada que puede ayudar a informar la investigación, como pronto veremos. Parte 2: Datos y modelo Preparación de datos de entrenamiento Estas son algunas de las características objetivo que buscamos con la generación de datos: Un número igual de cuerdas balanceadas y no balanceadas. Las cadenas tendrán una longitud par, ya que una cadena de longitud impar obviamente está desequilibrada; lo cual no sería una heurística muy interesante de aprender para el modelo. Todas las longitudes de cuerda (2-40) deberían ser igualmente probables. Para una longitud de cadena determinada, todas las posibles profundidades de anidamiento de paréntesis deberían ser igualmente probables. Un tema común es evidente: estamos tratando de hacer que todas las estadísticas de distribución imaginables tengan la misma probabilidad de reducir el sesgo en cualquier dirección dada, garantizar la solidez y negar las heurísticas obvias de ganancia rápida como una opción para el modelo. Para generar casos de falla, primero generaremos paréntesis válidos con las garantías mencionadas anteriormente y luego mutaremos la mitad de ellos para desequilibrarlos. 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 Esquema de generación de datos n.° 1: paseo aleatorio El primer intento de generación de paréntesis simplemente realiza un recorrido aleatorio. Pero como puede ver en los gráficos a continuación, el subespacio de paréntesis no balanceados es mucho mayor que el de los paréntesis balanceados; entonces tendremos que introducir la estocasticidad de manera diferente. 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() Esquema de generación de datos n.° 2: secuencia de anidamiento aleatoria codiciosa Podemos dividir la construcción de una cadena de paréntesis equilibrada en unidades discretas de paréntesis anidados. Para esta codiciosa construcción, en cada paso del proceso de generación de una cuerda se elige una profundidad de anidación de una cesta de profundidades viables (para respetar la longitud de la cuerda objetivo). Por ejemplo, para la longitud objetivo , son posibles las siguientes descomposiciones de anidamiento únicas: 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 Obtenga padres equilibrados 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] Veamos las frecuencias de las profundidades de anidación. 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() Y ahora, a ver las frecuencias de longitud. paren_len_hist = px.histogram(len_evals, title="Frequency of string lengths") paren_len_hist.show() Nota de generación de datos Tenga en cuenta que existe una tensión entre las siguientes propiedades potenciales de nuestra distribución de datos. Cada longitud de cuerda es igualmente probable. Cada subcadena de profundidad de anidamiento tiene la misma probabilidad en todas las cadenas. Esto se debe a que las subsecuencias de anidación baja tendrán más oportunidades de aparecer en una secuencia de anidación aleatoria determinada, como se muestra en los gráficos anteriores. Para contrarrestar esta tendencia natural de la secuencia puramente aleatoria, al generar una subcadena determinada de paréntesis, podríamos tomar muestras de una distribución sesgada para hacer que los valores de nido más profundos sean más probables. Esto se revisará después de una primera fase de formación. px.histogram(random_seq_parens, title="Frequency of balanced Parentheses").show() Crear paréntesis desequilibrados Nuestro conjunto de datos no puede tener solo paréntesis equilibrados. De modo que podemos crear una estrategia de generación de datos para derivar cadenas desequilibradas de nuestro conjunto de datos equilibrado. 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('((()))')) Obtenga un conjunto de datos de padres desequilibrados unbal_random_seq_parens = [make_parens_unbalanced(paren) for paren in random_seq_parens] Entrenamiento modelo Ahora que tenemos nuestros conjuntos de datos, por diversión, escribiremos nuestra arquitectura Transformer desde cero. Primero algunas configuraciones @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 Luego nuestro tokenizador para analizar entradas: 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) (Des)incrustaciones 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 Norma de capa práctica 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 Y por último ¡Atención! 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) Capas 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) Poniéndolo junto en un transformador 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) Utilidades de entrenamiento 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 Configuración de entrenamiento 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) Entrenamiento real 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}') En la Parte 3, investigaremos los aspectos internos de esta red entrenada. Haremos esto observando los patrones de atención y aplicando algunas de las herramientas de diagnóstico de interpretabilidad mecanicista, como los parches de activación, para construir un modelo mecanicista que comprenda cómo la red ha resuelto esta tarea. ¡Gracias por leer hasta aquí y nos vemos pronto en la Parte 3!
