Prevendo sequências simples com um transformer artesanal, sem treinamento

• Um experimento que constrói um transformer somente decodificador parecido com o GPT-2, projetando manualmente apenas os pesos, sem treinamento, para prever o padrão (aab)*
• A previsão de "aabaabaabaab..." precisa olhar para os dois tokens anteriores, então é melhor para revelar o funcionamento da atenção do que um padrão alternado simples
• O modelo é pequeno, com N_CTX=5, N_VOCAB=2 e N_EMBED=8, usando tokenização a=0, b=1 e embeddings one-hot
• Um único head de atenção distribui 0,5 de atenção para cada um dos dois tokens mais recentes, e calcula o próximo token usando o cancelamento da codificação a=1, b=-1
• Em contextos não ambíguos, mostrou 100,0% (27/27) de acerto, mas precisa de cerca de 4.000 FLOPs para uma única previsão com contexto de 5 tokens, sendo muito menos eficiente do que uma regra implementada diretamente

---

Um mini GPT-2 com pesos definidos à mão, sem treinamento

• O objetivo é entender de forma intuitiva o que cada componente do transformer e da atenção realmente faz
• O modelo não é treinado nem usa pesos pré-treinados; ele é construído em uma noite, definindo manualmente cada peso
• A estrutura é um transformer somente decodificador parecido com o GPT-2, e a implementação foi simplificada a partir da implementação picoGPT de jaymody
  • remoção do layer norm
  • uso de um único head em vez de multi-head attention
  • remoção da camada feed-forward mlp do bloco transformer

Por que escolher a sequência (aab)*

• No começo, a ideia era prever uma sequência como "ababababab", mas isso ficou fácil demais porque o transformer prevê uma sequência deslocada
  • Bastaria prever b se fosse a, e a caso contrário, sem precisar usar embedding posicional
• A tarefa final foi definida como prever "aabaabaabaab...", ou seja, a sequência (aab)*
  • se os dois tokens anteriores forem ab ou ba, o próximo token é a
  • se os dois tokens anteriores forem aa, o próximo token é b
  • bb é tratado como um caso fora do escopo da tarefa
• A tokenização é simples e lida com apenas dois símbolos
  • a é 0
  • b é 1

Dimensões do modelo e fluxo de cálculo

• Foram escolhidos três parâmetros para o modelo
  • N_CTX = 5: tamanho máximo de contexto que o modelo vê de uma vez
  • N_VOCAB = 2: os dois tokens a e b
  • N_EMBED = 8: tamanho do embedding que guarda espaço para token, posição e cálculo
• Embora a tarefa real só precise dos dois tokens anteriores, usar N_CTX=5 inclui a situação em que é preciso ignorar tokens irrelevantes
• A função gpt funciona na seguinte ordem
  • soma embeddings de token e posição com wte[inputs] + wpe[range(len(inputs))]
  • passa por um único bloco transformer
  • ao final, cria os logits no espaço do vocabulário com x @ wte.T

Embeddings com posição e token em one-hot

• wpe é o embedding posicional, usando as primeiras 5 dimensões do embedding como one-hot de posição
  • posição 0 é [1, 0, 0, 0, 0]
  • posição 4 é [0, 0, 0, 0, 1]
• wte é o embedding de token, usando as 2 dimensões seguintes como one-hot de token
  • o token a é [1, 0] nessas dimensões
  • o token b é [0, 1]
• A 8ª posição do embedding não é usada no início e vira um scratch space dentro do bloco transformer
• Por exemplo, "aabaa" é representado como uma matriz de embeddings 5 x 8, em que cada linha contém tanto o one-hot de posição quanto o one-hot de token

Como o head de atenção escolhe os dois tokens mais recentes

• O bloco transformer é composto por um único head de atenção e por uma camada linear c_proj que projeta o resultado da atenção de volta ao espaço de embedding
• c_attn é uma camada linear de tamanho embed_size x (embed_size * 3), que transforma o embedding de entrada em uma matriz qkv, depois dividida em q, k e v
• k separa o embedding posicional para representar a informação de posição de cada token
• q representa a faixa de posições que cada posição quer buscar, e q @ k.T produz a matriz de scores de atenção
• Depois do softmax e da causal mask, a matriz de atenção tem as seguintes propriedades
  • a primeira linha dá 100% de atenção apenas ao primeiro token
  • as linhas seguintes dão 0,5 de atenção para cada um dos dois tokens mais recentes acessíveis
• A causal mask impede olhar para tokens futuros somando um valor muito pequeno nas posições futuras; no código real, algo como -1e10
  • este modelo artesanal não foi projetado para espionar o futuro, mas a máscara foi mantida para ficar mais próximo da estrutura do GPT-2
• O escalonamento por np.sqrt(q.shape[-1]) ajuda a melhorar gradientes no treinamento real, mas não afeta este modelo feito à mão

Codificação em v e previsão por soma com cancelamento

• v transforma o one-hot do token em uma codificação a=1, b=-1
• Como o resultado da atenção faz a média dos dois tokens mais recentes com peso 0,5 para cada um, essa codificação calcula a seguinte regra
  • a, b → 0.5 * 1 + 0.5 * (-1) = 0
  • b, a → 0.5 * (-1) + 0.5 * 1 = 0
  • a, a → 0.5 * 1 + 0.5 * 1 = 1
• Como resultado, a 7ª posição da linha passa a conter o seguinte valor
  • 0 quando se deve prever a
  • 1 quando se deve prever b
• Na entrada "aabaa", a primeira linha pode gerar uma previsão b por falta de informação, mas as previsões seguintes batem com a regra de (aab)*

Enviando a previsão de volta ao espaço do vocabulário

• c_proj converte o valor da 7ª posição do resultado da atenção de volta para o formato one-hot do token
• Em vez de simplesmente criar [..., 1, 0, ...] ou [..., 0, 1, ...], ele cria um one-hot escalado por 1024
  • embedding[row, 5] = 1024 + (-1024) * prediction
  • embedding[row, 6] = 0 + 1024 * prediction
• O bloco transformer tem uma residual connection, então x = x + causal_self_attention(...) soma o embedding original de volta
• Como o sinal residual sobraria sem necessidade, usa-se a escala 1024 para sobrepujá-lo
• Por fim, calcula-se x @ wte.T para gerar os logits e aplica-se softmax
  • no contexto "aabaa", a linha final da previsão aponta para b
  • no treinamento, as previsões de todas as linhas são úteis, mas na inferência só a última linha é necessária

Resultado da geração e acurácia

• A função complete passa os últimos até 5 tokens ao modelo e escolhe o próximo token com argmax na última linha do resultado do softmax
• Exemplos de geração
  • complete("a") → a :: baabaabaab
  • complete("ba") → ba :: abaabaabaa
  • complete("abaab") → abaab :: aabaabaaba
• Mesmo com entradas fora do escopo, às vezes o modelo recupera o padrão repetitivo
  • complete("ababa") → ababa :: abaabaabaa
  • complete("bbbbb") → bbbbb :: aabaabaaba
• No teste com "aab" * 10, avaliando apenas contextos não ambíguos, a acurácia foi de 100,0% (27/27)

A diferença entre 4.000 FLOPs e 8 instruções

• Usando o contexto completo de 5 tokens, este modelo precisa de cerca de 4.000 operações de ponto flutuante para prever um único token
  • a maior parte é gasta no cálculo da atenção
  • isso pode ser reduzido com janela de contexto menor, fused multiply-add, cache de kv e outras técnicas
  • ainda assim, prever um único token exigiria centenas de instruções de máquina
• A mesma regra (aab)*, escrita diretamente em assembly x64, calcula o próximo token em 8 instruções
• Fica a pergunta se é possível treinar um modelo de linguagem 1000 vezes mais eficiente do que os modelos atuais para geração de linguagem natural

[GN⁺]

Comentários do Hacker News

• Há um trabalho relacionado chamado "Thinking Like Transformers"
  Ele apresenta uma linguagem de programação primitiva chamada RASP, composta por operações que podem ser modeladas com componentes de Transformer, e mostra que é possível escrever programas como histogramas ou ordenação
  Também há excelentes posts de blog de Sasha Rush e Gail Weiss, e pesquisas posteriores mostraram que programas do tipo RASP podem ser compilados diretamente para pesos de modelo reais, sem treinamento
  [1] https://arxiv.org/abs/2106.06981
  [2] https://srush.github.io/raspy/
  [3] https://arxiv.org/abs/2301.05062

  • Gosto muito da família RASP
    Se essa área parecer interessante, talvez valha dar uma olhada também no meu trabalho HandCrafted Transformers, em que escolhi manualmente os pesos de um modelo Transformer para fazê-lo realizar adição longa, de uma forma parecida com a que as pessoas aprendem no ensino fundamental
    [1] https://colab.research.google.com/github/newhouseb/handcraft...
  • Parece que uma linguagem funcional como Haskell se encaixaria bem nesse tipo de trabalho
    A direção oposta, de redes neurais para código, também parece muito interessante do ponto de vista da explicabilidade

• Eu achava que entendia bem Transformers, mas nunca tinha implementado um por conta própria
  Um dia implementei, e ele não funcionava nem treinava tão bem quanto o Transformer padrão do PyTorch; no fim, percebi que o motivo era eu ter ignorado o dropout
  Eu estava treinando adição de números e nunca mostrava o mesmo par duas vezes, então achava que overfitting seria impossível, mas o papel do dropout foi muito maior do que eu esperava
  Em resumo, é bom simplesmente implementar um Transformer por conta própria, e quanto mais do zero, melhor
  Todos que fizeram isso aprenderam algo inesperado, e cada pessoa percebeu coisas diferentes, desde paralelizar o treinamento em nível de tokens até como a retropropagação realmente funciona

  • Fico curioso se há alguma referência que ajude a abordar esse trabalho

• Também gosto dos materiais do Karpathy, mas o que finalmente fez Transformer fazer sentido para mim foi este vídeo: https://youtu.be/kWLed8o5M2Y?si=SJT5_lCJ0hSR7Z_k

• Venho pensando em algo parecido há algum tempo
  Será que não dá para criar uma interface intuitiva para os pesos de um modelo, que um especialista de domínio possa ajustar manualmente para acelerar o treinamento?
  Por exemplo, em um modelo de visão, ao detectar cones de trânsito, aumentar um conjunto de pesos correspondente a "alaranjadez"
  Assim, em vez de exigir milhares ou milhões de exemplos adicionais para calibrar corretamente a "alaranjadez", uma pessoa poderia acelerar o processo
  Claro que a dificuldade é que essa interface teria de ser mapeada para conjuntos de pesos com significados diferentes, e fico curioso se há algum motivo técnico para isso ser impossível

  • "Um especialista de domínio ajustar manualmente os pesos do modelo" soa parecido com a forma de reconhecimento de imagens antes do deep learning
    [1] https://www.youtube.com/watch?v=8SF_h3xF3cE&t=1358s
  • O motivo que você procura é chamado de The Bitter Lesson
    Em poucas palavras, formas de humanos ajudarem a IA quase sempre têm pior custo-benefício do que usar mais poder computacional
    Enquanto uma pessoa calibra uma camada de pesos para detectar cones de trânsito laranja, um cluster de GPUs treina a IA para detectar cones de trânsito, semáforos, árvores, outros carros e até cones de trânsito em tons de laranja ligeiramente diferentes
  • O número de camadas e pesos não está em uma escala que uma pessoa possa atualizar manualmente, e, mesmo que estivesse, é difícil demais gerenciar os efeitos a jusante de mudar os pesos
    Mesmo que você ajuste a imagem para enxergar melhor o laranja, se não puder monitorar ao mesmo tempo a precisão em todas as outras cores, é provável que crie problemas em outras cores sem perceber
  • O motivo técnico pelo qual isso é impossível, ou muito difícil, é que os pesos geralmente são muito difíceis de interpretar
    Não é como se um cluster específico de neurônios correspondesse a um conceito específico; em geral, tudo faz um pouco de tudo
  • O mecanismo de atenção dos Transformers não parece ser facilmente mapeável para semânticas compreensíveis por humanos
    Há parâmetros demais envolvidos

• O artigo dos Transformers é técnico demais, então eu sempre quis entendê-lo, mesmo que superficialmente, mas era difícil
  Este texto ajudou muito a entender como eles funcionam, e pelo menos o exemplo ficou muito claro
  Graças a ele, também consegui relembrar as matrizes que aprendi na faculdade

• Não seria uma espécie de máquina abstrata, como uma máquina de Turing ou uma máquina que faz parsing de expressões regulares?

  • Simplificando um pouco, é uma "máquina" que mapeia um conjunto de entradas para um conjunto de probabilidades da próxima saída
    Primeiro você define uma lista de tokens; por exemplo, para facilitar, digamos que sejam 24 caracteres
    Essa máquina recebe uma sequência de tokens de entrada, executa operações matriciais determinísticas e então gera uma lista de probabilidades para todos os tokens
    "Aprendizado" é apenas o processo de definir alguns dos números dentro das matrizes usadas nessas operações
    Vale notar que no código final há apenas um if, e ele serve para avaliar a precisão do resultado
    Toda a "lógica" vem dos resultados das operações matriciais
  • É bem difícil interpretar essas coisas como autômatos no sentido em que normalmente pensamos neles
    Em redes neurais, tudo em geral é meio nebuloso, e quase não há coisas como if/else, embora existam casos em que valores são "mascarados" como 0 ou -∞, como no exemplo de Transformer
    A saída também é quase sempre um conjunto de pontuações ou probabilidades; então, se um modelo que distingue fotos de cães e gatos produz algo como dog:0.95 cat:0.05, dizemos que ele previu cachorro porque a pontuação de cachorro é maior
    O mecanismo de atenção, que é o núcleo do Transformer, se baseia em uma espécie de operação de consulta suave
    Num sistema não nebuloso, você percorreria cada token da sequência, verificaria se ele é relevante para o token atual e, se fosse, faria alguma ação; mas, no Transformer, a relevância não é uma decisão binária
    Em vez disso, ele calcula uma pontuação contínua de relevância entre todos os pares de tokens da sequência e usa essa pontuação para decidir o próximo passo
    Só que algumas coisas não são fáceis de generalizar diretamente a partir de sistemas baseados em decisões binárias
    Por exemplo, essas pontuações de relevância são usadas como pesos para calcular uma média ponderada dos tokens do vocabulário, obtendo um "token médio" para a posição atual
    Não parece haver uma maneira fácil de interpretar isso como uma extensão de um processo baseado em lógica de ramificação
  • Não é como se o AllSpark tivesse tocado numa pilha de álgebra linear?
  • Isso mesmo
    Vale ver este artigo, que explica que Linear Transformers são, na verdade, Fast Weight Programmers: https://arxiv.org/abs/2102.11174
  • Redes neurais são máquinas de Turing
    Se você definir os pesos cuidadosamente, consegue fazê-las executar qualquer computação
    Só seria bom ter um compilador que não dependesse de aproximações

• Fico curioso sobre para que serviria a frase "talvez você fique com vontade de criar seu próprio modelo", além de ser um exercício de aprendizado para satisfazer a curiosidade
  Estou começando a sentir que modelos de machine learning complexos são impraticáveis para alguém lendo blogs em casa

  • No nanoGPT, se você pré-treina um modelo com Shakespeare, em 3 minutos ele chega a um nível de fidelidade ao material original comparável ao Jabberwocky, de Lewis Carroll
    Ele cria muitas palavras em inglês antigo que parecem plausíveis e aprende o básico da gramática inglesa e o formato de peças teatrais, entre outras coisas
    Achei bem surpreendente chegar a esse ponto em tão pouco tempo
    Treinar vários modelos localmente até um nível de fidelidade tipo Shakespeare-from-Wish.com poderia ajudar a julgar se você encontrou uma boa arquitetura e se chegou a hora de tentar escalá-la
  • O objetivo está no primeiro parágrafo do texto
    Ele diz que queria entender melhor Transformers e atenção, e que tinha lido The Illustrated Transformer, mas não conseguia formar uma intuição sobre o que as várias partes da atenção realmente faziam
    Algo como a diferença entre q e k, sem falar em v
  • É um excelente exercício de aprendizado
    Vai além de simplesmente satisfazer a curiosidade: ajuda a construir e aprofundar o entendimento
  • Talvez a pessoa simplesmente goste de hackear esse tipo de projeto? Estranho, mas possível

• Seria bom se o título pudesse incluir uma expressão como neural network
  Isto é sobre a arquitetura "Transformer" de machine learning, não sobre um conjunto de bobinas que acopla eletromagneticamente dois circuitos