Como usar a LIT para analisar modelos Gemma no Keras

Ver sobre a IA generativa Executar no Google Colab Consulte o código-fonte no GitHub Aprender em Codelabs

Introdução

Os produtos de IA generativa são relativamente novos, e os comportamentos de um aplicativo podem variar mais do que formas anteriores de software. Isso torna importante a sondagem dos modelos de machine learning que estão sendo usados, a análise de exemplos de comportamento do modelo e a investigação de surpresas.

A ferramenta de aprendizado de interpretabilidade (LIT; site, GitHub) é uma plataforma para depurar e analisar modelos de ML para entender por que e como eles se comportam dessa maneira.

Neste codelab, você aprenderá a usar a LIT para aproveitar ao máximo o modelo Gemma do Google. Este codelab demonstra como usar a saliência de sequências, uma técnica de interpretabilidade, para analisar diferentes abordagens de engenharia de comandos.

Objetivos de aprendizado:

  1. Entender a saliência de sequências e os usos dela na análise de modelos.
  2. Como configurar a LIT do Gemma para calcular as saídas de comandos e a saliência de sequência.
  3. Usando a saliência de sequências com o módulo LM Salience para entender o impacto dos designs de comandos nos resultados do modelo.
  4. Testar possíveis melhorias de comandos na LIT e conferir o impacto delas.

Observação: este codelab usa a implementação do KerasNLP (link em inglês) do Gemma e o TensorFlow v2 para o back-end. É altamente recomendável usar um kernel da GPU para acompanhar.

A saliência de sequências e os respectivos usos na análise de modelos

Os modelos generativos de texto para texto, como o Gemma, usam uma sequência de entrada na forma de texto tokenizado e geram novos tokens que são acompanhamentos ou conclusões típicos dessa entrada. Essa geração acontece um token por vez, anexando (em um loop) cada token recém-gerado à entrada, além de qualquer geração anterior, até que o modelo atinja uma condição de parada. Exemplos incluem quando o modelo gera um token de fim de sequência (EOS, na sigla em inglês) ou atinge o tamanho máximo predefinido.

Os métodos de saliência são uma classe de técnicas de Explainable AI (XAI) que pode informar quais partes de uma entrada são importantes para o modelo em diferentes partes da saída. A LIT é compatível com métodos de saliência para várias tarefas de classificação, que explicam o impacto de uma sequência de tokens de entrada no rótulo previsto. A saliência de sequência generaliza esses métodos para modelos generativos de texto para texto e explica o impacto dos tokens anteriores nos tokens gerados.

Você usará o método Grad L2 Norm para saliência de sequência, que analisa os gradientes do modelo e fornece uma magnitude da influência que cada token anterior tem na saída. Esse método é simples e eficiente, além de ter demonstrado bom desempenho em classificação e em outras configurações. Quanto maior a pontuação de saliência, maior a influência. Esse método é usado na LIT porque é bem compreendido e utilizado amplamente em toda a comunidade de pesquisa de interpretabilidade.

Métodos mais avançados de saliência baseados em gradiente incluem Grad ⋅ Entrada e gradientes integrados. Há também métodos baseados em ablação, como LIME e SHAP, que podem ser mais robustos, mas significativamente mais caros de calcular. Consulte este artigo para uma comparação detalhada dos diferentes métodos de saliência.

Saiba mais sobre a ciência dos métodos de saliência nesta introdução interativa introdutória sobre saliência.

Importações, ambiente e outros códigos de configuração

ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
bigframes 0.21.0 requires scikit-learn>=1.2.2, but you have scikit-learn 1.0.2 which is incompatible.
google-colab 1.0.0 requires ipython==7.34.0, but you have ipython 8.14.0 which is incompatible.

Eles podem ser ignorados.

Instalar a LIT e o PLN do Keras

Para este codelab, você precisará de uma versão recente de keras (3) keras-nlp (0.8.0) e lit-nlp (1.1), além de uma conta do Kaggle para fazer o download do modelo base.

pip install -q -U lit-nlp
pip uninstall -y umap-learn
pip install -q -U keras-nlp
pip install -q -U keras

Acesso ao Kaggle

Para fazer login no Kaggle, armazene o arquivo de credenciais do kaggle.json em ~/.kaggle/kaggle.json ou execute o seguinte em um ambiente do Colab. Consulte a documentação do pacote kagglehub para mais detalhes.

import kagglehub

kagglehub.login()

Aceite o contrato de licença da Gemma.

Como configurar a LIT para o Gemma

Como configurar modelos LIT

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import keras
import keras_nlp

# Run at half precision.
keras.config.set_floatx("bfloat16")

model_name = 'gemma_instruct_2b_en'
gemma_model = keras_nlp.models.GemmaCausalLM.from_preset(model_name)

O código a seguir inicializa os wrappers LIT para oferecer compatibilidade com a saliência no modelo Gemma. O framework LIT se refere a eles como modelos, mas, nesse caso, eles são apenas endpoints diferentes para a mesma gemma_model subjacente que você carregou acima. Isso permite que a LIT calcule gerações, tokenização e saliência sob demanda.

from lit_nlp.examples.models import instrumented_keras_lms

batch_size = 1
max_sequence_length = 512
init_models = instrumented_keras_lms.initialize_model_group_for_salience
models = init_models(model_name, gemma_model,
                     batch_size=batch_size,
                     max_length=max_sequence_length)

Como configurar conjuntos de dados da LIT

O Gemma é um modelo generativo de texto para texto que recebe uma entrada de texto e gera uma saída de texto. Os modelos da LIT pressupõem que os conjuntos de dados fornecerão os seguintes campos para dar suporte à geração:

  • prompt: a entrada para um KerasGenerationModel.
  • target: uma sequência de destino opcional, como uma resposta de "informações empíricas" (ouro) ou uma resposta pré-gerada do modelo.

A LIT inclui um pequeno conjunto de sample_prompts com exemplos de algumas fontes diferentes, como:

  • [GSM8K][GSM8K]: Resolução de problemas de matemática do ensino fundamental com exemplos few-shot.
  • [Gigaword Benchmark][gigaword]: geração de títulos para uma coleção de artigos curtos.
  • [Instrução constitucional][constitucional-prompting]: geração de novas ideias sobre como usar objetos com diretrizes/limites

Também é possível carregar facilmente seus próprios dados, seja como um arquivo .jsonl contendo registros com os campos prompt e opcionalmente target ([exemplo][jsonl-example]) ou de qualquer formato usando a API Dataset da LIT.

Execute a célula abaixo para carregar os comandos de amostra.

from lit_nlp.examples.datasets import lm as lm_data

datasets = {
  'sample_prompts': lm_data.PromptExamples(
      lm_data.PromptExamples.SAMPLE_DATA_PATH
  ),
}

Como configurar a interface da LIT

A LIT é uma ferramenta interativa de compreensão de modelos. Ela permite avaliações e sondagem de comportamento do modelo com intervenção humana. A interface da LIT facilita essa interação, permitindo que você:

  • visualizar os conjuntos de dados e as saídas do modelo em tempo real,
  • executar métodos de saliência para entender os tokens de entrada que impulsionam o comportamento do modelo
  • criar contrafatos para testar hipóteses.

A LIT permite tudo isso na mesma interface, reduzindo o atrito da alternância entre diferentes ferramentas. Isso é particularmente útil para tarefas como a engenharia de comandos, em que você vai se concentrar mais adiante neste codelab.

Esse layout de interface pode ser usado para qualquer outro modelo de linguagem generativa. Se você tem interesse em outros recursos além dos listados aqui, consulte a lista completa neste link.

from lit_nlp.api import layout
modules = layout.LitModuleName

LM_SALIENCE_LAYOUT = layout.LitCanonicalLayout(
    left={
        'Data Table': [modules.DataTableModule],
        'Datapoint Editor': [modules.DatapointEditorModule],
    },
    upper={  # if 'lower' not specified, this fills the right side
        'Salience': [modules.LMSalienceModule],
    },
    layoutSettings=layout.LayoutSettings(leftWidth=40),
    description='Custom layout for language model salience.',
)

Essa célula inicializa o servidor da LIT. Isso pode levar alguns segundos, já que também executa o modelo nos comandos de amostra e armazena o resultado em cache.

from lit_nlp import notebook as lit_notebook

lit_widget = lit_notebook.LitWidget(
    models=models,
    datasets=datasets,
    layouts={'default': LM_SALIENCE_LAYOUT},
    default_layout='default',
)

Agora você pode mostrar a interface:

lit_widget.render(height=800)

<IPython.core.display.Javascript object>

Também é possível abrir a LIT como página inteira em uma nova guia. Execute esta célula e clique no link exibido:

lit_widget.render(open_in_new_tab=True)

<IPython.core.display.Javascript object>

Como analisar prompts de few shot para o Gemma na LIT

Hoje, a criação de comandos é tanto arte quanto ciência, e a LIT pode ajudar você a melhorar empiricamente os comandos para modelos de linguagem grandes, como o Gemma. A seguir, você verá um exemplo de como a LIT pode ser usada para analisar os comportamentos da Gemma, prever possíveis problemas e melhorar a segurança dela.

Identificar erros em comandos complexos

Duas das técnicas de comandos mais importantes para protótipos e aplicativos de alta qualidade baseados em LLMs são os comandos few-shot (incluindo exemplos do comportamento desejado no comando) e a cadeia de pensamento (incluindo uma explicação ou raciocínio antes da saída final do LLM). Mas criar um comando eficaz ainda é desafiador.

Considere ajudar uma pessoa a avaliar se ela gostará de comida com base nos gostos dela. Um protótipo inicial de modelo de comando de cadeia de pensamento pode ter esta aparência:

def analyze_menu_item_template(food_likes, food_dislikes, menu_item):
  return f"""Analyze a menu item in a restaurant.

## For example:

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Onion soup
Analysis: it has cooked onions in it, which you don't like.
Recommendation: You have to try it.

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Baguette maison au levain
Analysis: Home-made leaven bread in france is usually great
Recommendation: Likely good.

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Macaron in france
Analysis: Sweet with many kinds of flavours
Recommendation: You have to try it.

## Now analyze one more example:

Taste-likes: {food_likes}
Taste-dislikes: {food_dislikes}
Suggestion: {menu_item}
Analysis:"""

Você identificou os problemas com o comando? A LIT ajuda você a examinar o comando usando o módulo Salience do LM.

Usar a saliência de sequências para depuração

Este módulo destaca partes do comando que o modelo atende ao gerar a resposta. A saliência é calculada no menor nível possível (ou seja, para cada token de entrada), mas a LIT pode agregar a saliência de tokens em períodos maiores e mais interpretáveis, como linhas, frases ou palavras. Saiba mais sobre saliência e como usar para identificar vieses não intencionais em nossa Saliência explorável.

Vamos começar inserindo um novo exemplo de entrada para as variáveis do modelo de comando:

food_likes = """Cheese"""
food_dislikes = """Can't eat eggs"""
menu_item = """Quiche Lorraine"""

prompt = analyze_menu_item_template(food_likes, food_dislikes, menu_item)
print(prompt)

fewshot_mistake_example = {'prompt': prompt}  # you'll use this below

Analyze a menu item in a restaurant.

## For example:

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Onion soup
Analysis: it has cooked onions in it, which you don't like.
Recommendation: You have to try it.

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Baguette maison au levain
Analysis: Home-made leaven bread in france is usually great
Recommendation: Likely good.

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Macaron in france
Analysis: Sweet with many kinds of flavours
Recommendation: You have to try it.

## Now analyze one more example:

Taste-likes: Cheese
Taste-dislikes: Can't eat eggs
Suggestion: Quiche Lorraine
Analysis:

Se a interface da LIT estiver aberta na célula acima ou em uma guia separada, use o editor de ponto de dados da LIT para adicionar este comando:

1_editor_de_dados.png

Outra maneira é renderizar novamente o widget diretamente com o comando de interesse:

lit_widget.render(data=[fewshot_mistake_example])

<IPython.core.display.Javascript object>

Observe a conclusão surpreendente do modelo:

Taste-likes: Cheese
Taste-dislikes: Can't eat eggs
Suggestion: Quiche Lorraine
Analysis: A savoury tart with cheese and eggs
Recommendation: You might not like it, but it's worth trying.

Por que o modelo está sugerindo que você coma algo que você claramente disse que não pode comer?

A saliência de sequência pode ajudar a destacar o problema raiz, que está nos nossos exemplos de few-shot. No primeiro exemplo, o raciocínio da cadeia de pensamento na seção de análise it has cooked onions in it, which you don't like não corresponde à recomendação final You have to try it.

No módulo "Salience" do LM, selecione "Frases" e a linha de recomendação. A interface será parecida com esta:

3_few_shots_mistake..png

Agora, vamos corrigir a "Recomendação" no primeiro exemplo para Avoid e tentar de novo. A LIT tem este exemplo pré-carregado nos prompts de amostra. Assim, você pode usar esta pequena função utilitária para capturá-lo:

def get_fewshot_example(source: str) -> str:
  for example in datasets['sample_prompts'].examples:
    if example['source'] == source:
      return example['prompt']
  raise ValueError(f'Source "{source}" not found in the dataset.')

lit_widget.render(data=[{'prompt': get_fewshot_example('fewshot-fixed')}])

<IPython.core.display.Javascript object>

A conclusão do modelo fica assim:

Taste-likes: Cheese
Taste-dislikes: Can't eat eggs
Suggestion: Quiche Lorraine
Analysis: This dish contains eggs and cheese, which you don't like.
Recommendation: Avoid.

Uma lição importante a ser retirada é: a prototipagem antecipada ajuda a revelar riscos que talvez você não pense com antecedência, e a natureza propensa a erros dos modelos de linguagem significa que é preciso projetar proativamente para os erros. Mais discussões sobre isso podem ser encontradas no Guia People + AI para projetar com IA.

Embora o comando de few shots corrigido seja melhor, ele ainda não está certo: ele diz ao usuário para evitar ovos, mas o raciocínio não está certo, diz que ele não gosta de ovos, quando, na verdade, o usuário declarou que não pode comer ovos. Na seção a seguir, você verá como melhorar.

Testar hipóteses para melhorar o comportamento do modelo

Com a LIT, você pode testar mudanças nos comandos na mesma interface. Nesta instância, você testará a adição de uma constituição para melhorar o comportamento do modelo. Constituições se referem a comandos de design com princípios para orientar a geração do modelo. Os métodos recentes permitem até mesmo a derivação interativa de princípios constitucionais.

Vamos usar essa ideia para melhorar ainda mais o comando. Adicione uma seção com os princípios para a geração na parte de cima do comando, que agora começa da seguinte maneira:

Analyze a menu item in a restaurant.

* The analysis should be brief and to the point.
* The analysis and recommendation should both be clear about the suitability for someone with a specified dietary restriction.

## For example:

Taste-likes: I've a sweet-tooth
Taste-dislikes: Don't like onions or garlic
Suggestion: Onion soup
Analysis: it has cooked onions in it, which you don't like.
Recommendation: Avoid.

...


lit_widget.render(data=[{'prompt': get_fewshot_example('fewshot-constitution')}])

<IPython.core.display.Javascript object>

Com essa atualização, o exemplo pode ser executado novamente e observar uma saída muito diferente:

Taste-likes: Cheese
Taste-dislikes: Can't eat eggs
Suggestion: Quiche Lorraine
Analysis: This dish containts eggs, which you can't eat.
Recommendation: Not suitable for you.

A saliência do comando pode ser reanalisada para ajudar a ter uma noção do motivo dessa mudança:

3_few_shot_constitution.png

Observe que a recomendação é muito mais segura. Além disso, a mensagem "Não adequado para você" é influenciado pelo princípio de indicar claramente a adequação de acordo com a restrição alimentar e a análise (a chamada cadeia de pensamento). Isso dá mais confiança de que a saída está acontecendo pelo motivo certo.

Incluir equipes não técnicas na sondagem e exploração de modelos

A interpretabilidade deve ser um esforço em equipe, que abrange conhecimentos sobre XAI, políticas, jurídicos e muito mais.

Interagir com modelos nos estágios iniciais de desenvolvimento sempre exigiu experiência técnica significativa, o que dificultava o acesso e a sondagem de alguns colaboradores. As ferramentas não existiam historicamente para que essas equipes participassem das fases iniciais de prototipagem.

Com a LIT, esperamos que esse paradigma possa mudar. Como você aprendeu neste codelab, o meio visual e a capacidade interativa da LIT de examinar a saliência e explorar exemplos podem ajudar diferentes partes interessadas a compartilhar e comunicar descobertas. Isso permite que você traga uma diversidade mais ampla de colegas de equipe para exploração, sondagem e depuração do modelo. A exposição a esses métodos técnicos pode melhorar o entendimento de como os modelos funcionam. Além disso, um conjunto mais diverso de experiências nos primeiros testes de modelo também pode ajudar a descobrir resultados indesejados que podem ser melhorados.

Recapitulação

Recapitulando:

  • A interface da LIT oferece uma interface para execução interativa de modelos, permitindo que os usuários gerem saídas diretamente e testem cenários "e se". Isso é particularmente útil para testar diferentes variações de comandos.
  • O módulo Salience do LM oferece uma representação visual da saliência e tem granularidade de dados controlável para que você possa se comunicar sobre construções centradas em pessoas (por exemplo, frases e palavras) em vez de em modelos (por exemplo, tokens).

Quando você encontrar exemplos problemáticos nas avaliações do modelo, coloque-os na LIT para depuração. Comece analisando a maior unidade sensata de conteúdo que você possa pensar que se relaciona logicamente com a tarefa de modelagem, use as visualizações para ver onde o modelo está atendendo corretamente ou incorretamente ao conteúdo do comando e, em seguida, detalhe os detalhes em unidades menores de conteúdo para descrever com mais detalhes o comportamento incorreto que você está vendo e identificar possíveis correções.

Por fim: o Lit está sempre melhorando. Saiba mais sobre nossos recursos e compartilhe suas sugestões aqui.

Apêndice: como a LIT calcula a saliência da sequência

A LIT calcula a saliência da sequência em um processo de várias etapas.

  1. Dada uma string de entrada (comando e a geração do modelo ou uma sequência de destino "ouro"), tokenize-a para entrada do modelo.
  2. Calcule uma sequência de "destino" rolando os tokens de entrada uma posição à esquerda.
  3. Extraia os embeddings e calcule a perda por token entre as sequências de geração e de destino.
  4. Mascarar a perda para isolar os tokens que você quer explicar.
  5. Use a função tf.GradientTape.gradient() (link em inglês) para calcular os gradientes dos embeddings de entrada em relação à perda mascarada.
  6. Processe os gradientes para dar uma única pontuação a cada token de entrada. Por exemplo, considerando a norma L2 do gradiente em cada posição.

Apêndice: calcular saliência de maneira programática

É possível calcular pontuações de saliência diretamente no Python usando as mesmas etapas acima em que a ferramenta LIT é executada internamente. Você fará isso em três etapas:

  1. Preparar um exemplo e executar o tokenizador de modelo
  2. Preparar uma máscara que selecione quais tokens (previstos) serão explicados
  3. Chame o wrapper de saliência.

Criar um exemplo de entrada para LIT

{'prompt': 'Keras is a',
 'target': ' deep learning library for Python that provides a wide range of tools and functionalities for building, training, and evaluating deep learning models.\n\n**'}

Uma observação sobre as convenções de chamada: o tokenizador e o wrapper de saliência usam a API Model da LIT, em que a função .predict() recebe uma lista de exemplos (dicts) e retorna um gerador de respostas (dicts). Isso é muito mais flexível ao trabalhar com conjuntos de dados maiores ou modelos mais lentos, mas significa que, se você só quer previsões em um exemplo, agora é necessário envolvê-lo com algo como: list(model.predict([example])[0]

Receber tokens para selecionar destinos de explicação

array(['<bos>', 'K', 'eras', '▁is', '▁a', '▁deep', '▁learning',
       '▁library', '▁for', '▁Python', '▁that', '▁provides', '▁a', '▁wide',
       '▁range', '▁of', '▁tools', '▁and', '▁functionalities', '▁for',
       '▁building', ',', '▁training', ',', '▁and', '▁evaluating', '▁deep',
       '▁learning', '▁models', '.', '\n\n', '**'], dtype='<U16')

Para calcular a saliência, é necessário criar uma máscara de destino que especifique quais tokens (previstos) serão explicados. A máscara de destino será uma matriz com o mesmo tamanho dos tokens, com 1s nas posições dos tokens que você quer explicar. Vamos usar ▁training e ▁evaluating como destino:

Preparar máscara de destino

{'prompt': 'Keras is a',
 'target': ' deep learning library for Python that provides a wide range of tools and functionalities for building, training, and evaluating deep learning models.\n\n**',
 'target_mask': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 1., 1., 1., 1., 0., 0., 0., 0., 0., 0.],
       dtype=float32)}

Chamar o modelo de saliência

{'grad_l2': array([45.75, 36.75, 61, 5.40625, 4.09375, 5.625, 6.46875, 7.3125, 3.375,
        5.03125, 3.23438, 4.5625, 2.375, 3.40625, 2.75, 1.97656, 3.95312,
        3.42188, 14.125, 4.53125, 11.375, 12.625, 18.5, 4.5625, 6.5, 0, 0,
        0, 0, 0, 0, 0], dtype=bfloat16),
 'grad_dot_input': array([-4.03125, 3.04688, -7.03125, -0.800781, 0.769531, -0.679688,
        -0.304688, 2.04688, 0.275391, -1.25781, -0.376953, -0.0664062,
        -0.0405273, -0.357422, 0.355469, -0.145508, -0.333984, 0.0181885,
        -5.0625, 0.235352, -0.470703, 2.25, 3.90625, -0.199219, 0.929688,
        0, 0, 0, 0, 0, 0, 0], dtype=bfloat16),
 'tokens': array(['<bos>', 'K', 'eras', '▁is', '▁a', '▁deep', '▁learning',
        '▁library', '▁for', '▁Python', '▁that', '▁provides', '▁a', '▁wide',
        '▁range', '▁of', '▁tools', '▁and', '▁functionalities', '▁for',
        '▁building', ',', '▁training', ',', '▁and', '▁evaluating', '▁deep',
        '▁learning', '▁models', '.', '\n\n', '**'], dtype='<U16')}

Pronto! As pontuações nos campos grad_l2 e grad_dot_input estão alinhadas a tokens e são as mesmas da interface LIT.

Observe que as últimas pontuações são 0: como o modelo é de linguagem da esquerda para a direita, os tokens à direita do período desejado não influenciam a previsão.