TD NLP - Transformers

Author

Augustin Chevallier

1 Un modèle de langage avec un transformer

Ce TD fait suite au TD NLP - Bases. On reprend le même dataset (TinyStories) et le même tokenizer caractère, mais on remplace le réseau feedforward par un transformer. On change ensuite de tokenizer.

1.1 Exercice 1 : remplacer le FFN par un transformer

Reprendre le chargement des données et le tokenizer caractère du TD précédent (fonctions encode et decode). On ajoute au vocabulaire deux tokens spéciaux, comme on l’avait fait pour <UNK> : <BOS> (beginning of sequence) et <EOS> (end of sequence), qui marqueront le début et la fin de chaque histoire.

from datasets import load_dataset

train = load_dataset("roneneldan/TinyStories", split="train[:1%]")
val = load_dataset("roneneldan/TinyStories", split="validation[:10%]")

textes_train = train["text"]
textes_val = val["text"]

1.1.1 Le dataset

Le FFN du TD précédent prenait 10 tokens et prédisait un seul token. Un transformer prend une séquence de \(T\) tokens et renvoie une séquence de \(T\) vecteurs : on peut donc prédire le token suivant à chaque position.

Un exemple est maintenant un couple de séquences de longueur \(w\), décalées de 1 :

\[ x = (x_i, \dots, x_{i+w-1}), \qquad y = (x_{i+1}, \dots, x_{i+w}) \]

Comme au TD précédent, on concatène toutes les histoires en un seul tenseur d’ids, mais en encadrant chaque histoire par <BOS> et <EOS> :

\[ \texttt{<BOS>}\; \text{histoire}_1 \;\texttt{<EOS>}\; \texttt{<BOS>}\; \text{histoire}_2 \;\texttt{<EOS>}\; \texttt{<BOS>}\; \text{histoire}_3 \;\dots \]

  1. Construire ce tenseur pour le train et la validation. Écrire une classe Dataset dont l’élément \(i\) est le couple (data[i:i+w], data[i+1:i+w+1]), avec \(w = 128\).
  2. Combien de prédictions fait-on pour un seul exemple ? Pourquoi est-ce plus efficace que le dataset du TD précédent ?
  3. La sortie en position \(t\) doit prédire \(x_{t+1}\) en ne regardant que \(x_1, \dots, x_t\). Quel mécanisme du cours permet de le garantir ?
  4. Une fenêtre peut contenir la fin d’une histoire, <EOS> <BOS>, puis le début de la suivante. Les tokens de la nouvelle histoire peuvent-ils porter attention à l’ancienne ? Est-ce grave ? Comment modifier le masque d’attention pour l’interdire ?

1.1.2 Le modèle

Étape Couche Taille de la sortie
entrée ids des tokens \((B, T)\)
embedding nn.Embedding(vocab_size, d) + nn.Embedding(w, d) (position) \((B, T, d)\)
transformer \(L\) couches nn.TransformerEncoderLayer, avec masque causal \((B, T, d)\)
normalisation nn.LayerNorm(d) \((B, T, d)\)
head nn.Linear(d, vocab_size) \((B, T, \text{vocab\_size})\)

On propose \(d = 128\), \(h = 4\) têtes, \(L = 4\) couches.

Indices : on empile les couches à la main, comme dans le cours.

# dans __init__
self.couches = nn.ModuleList([
    nn.TransformerEncoderLayer(d, h, dim_feedforward=4 * d, batch_first=True, norm_first=True)
    for _ in range(L)])
self.ln_finale = nn.LayerNorm(d)

# dans forward
mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), diagonal=1)
for couche in self.couches:
    e = couche(e, src_mask=mask)

Optionnel : couche(e, src_mask=mask, is_causal=True) indique à PyTorch que le masque est causal, ce qui lui permet d’utiliser une implémentation plus rapide de l’attention. Le masque doit quand même être fourni.

  1. L’encodage positionnel est ici appris : une deuxième table d’embedding, indexée par la position \(0, \dots, w-1\), ajoutée à l’embedding du token. Pourquoi en a-t-on besoin ?
  2. On utilise norm_first=True : de quelle variante du cours s’agit-il ? Pourquoi ajoute-t-on une LayerNorm finale avant la head ?
  3. Coder le modèle. Combien a-t-il de paramètres ? Comparer au FFN du TD précédent.

1.1.3 Entraînement

  1. Entraîner avec AdamW (lr=1e-3) et des batchs de 64, pendant quelques milliers d’itérations. Suivre la loss de validation et la comparer à celle du FFN.
  2. Sans masque. Relancer quelques centaines d’itérations en retirant le masque causal. Que devient la loss d’entraînement ? Pourquoi ce modèle est-il inutilisable pour générer du texte ?
  3. Loss par position. Calculer la loss de validation séparément pour chaque position \(t = 0, \dots, w-1\) de la fenêtre, et la tracer. Expliquer la forme de la courbe. La comparaison avec le FFN de la question 8 est-elle équitable ?

1.1.4 Génération

  1. Adapter la fonction generate(prompt, n, T) du TD précédent. Est-il encore nécessaire de compléter les prompts courts ? Pourquoi faut-il ne garder que les \(w\) derniers tokens comme contexte ?
  2. Générer à partir de "Once upon a time" pour les températures \(0\), \(0.5\) et \(1\). Comparer au FFN.
  3. Générer à partir de "Once upon a time, there was a girl named Mia.". Le nom du personnage réapparaît-il plus loin ? Pourquoi le FFN ne pouvait-il pas le faire ?
  4. Générer une histoire complète : partir d’un prompt réduit au seul token <BOS>, et s’arrêter quand le modèle génère <EOS>. Qu’apprend le modèle aux positions où l’entrée est <BOS> ? Et à celles où l’entrée est <EOS> ?

1.2 Exercice 2 : changer de tokenizer

Le tokenizer caractère donne des séquences très longues. On entraîne un tokenizer BPE (comme celui de GPT-2, mais avec un petit vocabulaire) sur nos histoires :

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2").train_new_from_iterator(textes_train, vocab_size=1000)

  1. Quelle est la taille du vocabulaire ? Afficher les tokens de la première phrase avec tokenizer.tokenize(...). Que représente le caractère Ġ ? Quels mots sont découpés en plusieurs tokens ?

  2. Combien de tokens contient le train ? Combien de caractères y a-t-il en moyenne par token ? Combien de caractères couvre maintenant une fenêtre de \(w = 128\) tokens ?

  3. Reprendre l’exercice 1 avec ce tokenizer, sans changer l’architecture.

    Le token de fin d’histoire est maintenant tokenizer.eos_token (<|endoftext|>), hérité de GPT-2. Le tokenizer de GPT-2 n’a pas de token de début, on l’ajoute :

    tokenizer.add_special_tokens({"bos_token": "<|startoftext|>"})

    On peut ensuite insérer tokenizer.bos_token et tokenizer.eos_token directement dans le texte avant de tokeniser :

    ids = tokenizer([tokenizer.bos_token + h + tokenizer.eos_token for h in textes_train])["input_ids"]
data_train = torch.tensor([i for histoire in ids for i in histoire])

    Quelles couches changent de taille ? Combien de paramètres a le modèle ?

  4. Comparer les loss. La loss par token des deux modèles est-elle comparable ? Pourquoi ? Pour comparer, convertir en loss par caractère : \[ \text{loss par caractère} = \text{loss par token} \times \frac{\text{nombre de tokens}}{\text{nombre de caractères}} \] Quel modèle est le meilleur ?

  5. Comparer les textes générés. Le modèle fait-il encore des fautes d’orthographe ?

  6. Recommencer avec vocab_size égal à 300 puis 4000. Quels sont les avantages et inconvénients d’un grand vocabulaire ? Pourquoi le tokenizer de GPT-2 (50 257 tokens) serait-il un mauvais choix ici ?

2 Pour aller plus loin

2.1 Exercice 3 : visualiser l’attention

Afficher la matrice d’attention de chaque tête pour une phrase courte. Les poids ne sont pas renvoyés par nn.TransformerEncoderLayer, mais on peut appeler directement son module d’attention :

x_norm = layer.norm1(x)   # avec norm_first=True, la LN est appliquée avant l'attention
_, poids = layer.self_attn(x_norm, x_norm, x_norm, attn_mask=mask,
                           need_weights=True, average_attn_weights=False)   # (B, h, T, T)

Retrouve-t-on des têtes qui regardent le token précédent, ou le début d’un mot ?

2.2 Exercice 4 : transformer from scratch

Remplacer nn.TransformerEncoderLayer par votre propre implémentation, y compris celle du mécanisme d’attention multi-tête. Vérifier qu’on obtient des loss similaires.