7 - Preentrenamiento y fine-tuning, en miniatura

En el capítulo 7, Ana enseña al equipo un iceberg: una masa enorme bajo el agua con la etiqueta preentrenamiento y una pequeña punta con la etiqueta fine-tuning. Casi nadie entrena un LLM desde cero para una tarea concreta. Primero el modelo se preentrena con una cantidad enorme de texto, prediciendo el siguiente token, y sale con un conocimiento general del lenguaje y del mundo, pero sin optimizar para nada más allá de completar texto. Después se ajusta (fine-tuning) para una tarea, con muchos menos datos y cómputo; uno de los tipos del libro es el ajuste de instrucciones, enseñar al modelo a responder preguntas o a seguir instrucciones.

En este cuaderno reproducimos las dos fases con un GPT diminuto, como el del cuaderno Un modelo de lenguaje tipo GPT:

1. Importar las bibliotecas.
2. El corpus de preentrenamiento: hojas y notas sobre 16 diagnósticos.
3. El preentrenamiento.
4. Lo que sabe el modelo preentrenado, y lo que no hace.
5. Los datos de fine-tuning: preguntas sobre 10 de los 16 diagnósticos.
6. El fine-tuning, con y sin texto de preentrenamiento.
7. Las mismas preguntas, entrenando desde cero.

Paso 1: Importar las bibliotecas

---

import copy
import math
import numpy as np
import torch
import torch.nn as nn

torch.manual_seed(0)

<torch._C.Generator at 0x111572730>

Paso 2: El corpus de preentrenamiento: hojas y notas sobre 16 diagnósticos

---

El texto de preentrenamiento es el tipo de texto que produce un hospital, sin preguntas ni respuestas: 4.000 hojas de alta, como las del cuaderno GPT, y 4.000 notas clínicas breves que cuentan los mismos hechos con otras palabras ("el reflujo se trata con omeprazol", "los pacientes con gota deben estar atentos a..."). Hay 16 diagnósticos, cada uno con su fármaco y su signo de alarma. Como en un LLM real, el vocabulario se fija antes del preentrenamiento y ya contiene las palabras de las preguntas del paso 5, aunque el texto de preentrenamiento nunca use algunas de ellas.

cases = {'neumonía': ('el antibiótico', 'dificultad para respirar'), 'migraña': ('ibuprofeno', 'visión borrosa'),
         'hipertensión': ('las pastillas de la tensión', 'dolor en el pecho'), 'gastroenteritis': ('suero oral', 'sangre en las heces'),
         'cistitis': ('fosfomicina', 'dolor de espalda'), 'asma': ('el inhalador', 'labios morados'), 'gota': ('colchicina', 'una articulación caliente e hinchada'),
         'insuficiencia cardiaca': ('furosemida', 'piernas hinchadas'), 'hipotiroidismo': ('levotiroxina', 'palpitaciones'),
         'diabetes': ('metformina', 'mucha sed'), 'reflujo': ('omeprazol', 'heces negras'), 'anemia': ('las pastillas de hierro', 'desmayos'),
         'herpes zóster': ('aciclovir', 'una erupción cerca del ojo'), 'trombosis': ('el anticoagulante', 'falta de aire repentina'),
         'conjuntivitis': ('el colirio', 'pérdida de visión'), 'otitis': ('amoxicilina', 'rigidez de cuello')}
frequencies = ['dos veces al día', 'cada ocho horas', 'una vez al día']
general_advice = ['beba mucha agua .', 'descanse en casa unos días .', 'no conduzca hoy .', 'haga comidas ligeras .',
                  'evite el alcohol .', 'puede volver al trabajo la semana que viene .', 'guarde esta hoja .']
note_templates = ['{dx} se trata con {drug} .', 'los pacientes con {dx} deben estar atentos a {warning} .',
                  '{drug} se usa para {dx} .', 'en {dx} el signo de alarma es {warning} .']
drug_questions = ['¿ qué debe tomar un paciente con {dx} ?', '¿ qué fármaco se usa para {dx} ?', '¿ cuál es el tratamiento de {dx} ?']
warning_questions = ['¿ cuándo debe volver un paciente con {dx} ?', '¿ qué signo de alarma importa en {dx} ?']
TAKE, QUESTION, ANSWER = 'tome', ['pregunta', ':'], ['respuesta', ':']

def discharge_sheet(rng):
    diagnosis = list(cases)[rng.integers(len(cases))]
    drug, warning = cases[diagnosis]
    words = ['diagnóstico', ':'] + diagnosis.split() + ['.']
    for advice in rng.choice(general_advice, rng.integers(0, 4), replace=False):
        words += advice.split()
    words += [TAKE] + drug.split() + rng.choice(frequencies).split() + ['.']
    for advice in rng.choice(general_advice, rng.integers(0, 3), replace=False):
        words += advice.split()
    return words + ['vuelva', 'si', 'tiene'] + warning.split() + ['.', '</s>'], diagnosis

def clinical_note(rng):
    diagnosis = list(cases)[rng.integers(len(cases))]
    drug, warning = cases[diagnosis]
    template = note_templates[rng.integers(len(note_templates))]
    return template.format(dx=diagnosis, drug=drug, warning=warning).split() + ['</s>']

def question_answer(diagnosis, template):
    drug, warning = cases[diagnosis]
    answer = drug if template in drug_questions else warning
    return QUESTION + template.format(dx=diagnosis).split() + ANSWER + answer.split() + ['.', '</s>'], answer

rng = np.random.default_rng(0)
documents = [discharge_sheet(rng)[0] for _ in range(4000)] + [clinical_note(rng) for _ in range(4000)]
documents = [documents[i] for i in np.random.default_rng(1).permutation(len(documents))]    # hojas y notas mezcladas

all_questions = [question_answer(d, t)[0] for d in cases for t in drug_questions + warning_questions]
vocab = sorted({w for doc in documents + all_questions for w in doc} | {'<s>'})
token_id = {w: i for i, w in enumerate(vocab)}
def to_stream(docs):
    return torch.tensor([token_id[w] for doc in docs for w in ['<s>'] + doc])
pretraining_stream = to_stream(documents)

for doc in documents[:4]:
    print(" ".join(doc))
print(f"\nDiagnósticos: {len(cases)}, vocabulario: {len(vocab)} tokens, corpus de preentrenamiento: {len(pretraining_stream):,} tokens")

los pacientes con asma deben estar atentos a labios morados . </s>
diagnóstico : hipotiroidismo . puede volver al trabajo la semana que viene . guarde esta hoja . evite el alcohol . tome levotiroxina dos veces al día . haga comidas ligeras . guarde esta hoja . vuelva si tiene palpitaciones . </s>
diagnóstico : hipertensión . guarde esta hoja . descanse en casa unos días . tome las pastillas de la tensión dos veces al día . evite el alcohol . beba mucha agua . vuelva si tiene dolor en el pecho . </s>
diagnóstico : neumonía . tome el antibiótico cada ocho horas . no conduzca hoy . vuelva si tiene dificultad para respirar . </s>

Diagnósticos: 16, vocabulario: 142 tokens, corpus de preentrenamiento: 175,864 tokens

Paso 3: El preentrenamiento

---

El modelo del cuaderno GPT, con el mismo objetivo: predecir el siguiente token en ventanas de 64 tokens. train toma lotes de un flujo de texto o, si le damos un segundo, la mitad de cada lote de cada uno; lo necesitaremos en el paso 6. El preentrenamiento es la fase larga: 3.000 pasos.

context_length = 64

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.heads = heads
        self.qkv, self.out = nn.Linear(d_model, 3 * d_model), nn.Linear(d_model, d_model)
        self.register_buffer('mask', torch.tril(torch.ones(context_length, context_length)).bool())

    def forward(self, x):
        batch, T, d = x.shape
        q, k, v = [m.view(batch, T, self.heads, d // self.heads).transpose(1, 2) for m in self.qkv(x).split(d, dim=2)]
        scores = (q @ k.transpose(-2, -1) / math.sqrt(d // self.heads)).masked_fill(~self.mask[:T, :T], float('-inf'))
        return self.out((scores.softmax(dim=-1) @ v).transpose(1, 2).reshape(batch, T, d))

class TransformerBlock(nn.Module):
    def __init__(self, d_model, heads):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.attention = CausalSelfAttention(d_model, heads)
        self.ffn = nn.Sequential(nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model))

    def forward(self, x):
        x = x + self.attention(self.norm1(x))
        return x + self.ffn(self.norm2(x))

class TinyGPT(nn.Module):
    def __init__(self, vocab_size, d_model=64, heads=4, layers=2):
        super().__init__()
        self.tokens, self.positions = nn.Embedding(vocab_size, d_model), nn.Embedding(context_length, d_model)
        self.blocks = nn.ModuleList([TransformerBlock(d_model, heads) for _ in range(layers)])
        self.norm = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.lm_head.weight = self.tokens.weight
        nn.init.normal_(self.tokens.weight, std=0.02)
        nn.init.normal_(self.positions.weight, std=0.02)

    def forward(self, x):
        h = self.tokens(x) + self.positions(torch.arange(x.shape[1]))
        for block in self.blocks:
            h = block(h)
        return self.lm_head(self.norm(h))

def windows(stream, n, generator):
    starts = torch.randint(0, len(stream) - context_length - 1, (n,), generator=generator)
    return (torch.stack([stream[s:s + context_length] for s in starts]),
            torch.stack([stream[s + 1:s + context_length + 1] for s in starts]))

def train(model, stream, steps, lr, mix_with=None, seed=0):
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    generator = torch.Generator().manual_seed(seed)
    for _ in range(steps):
        if mix_with is None:
            x, y = windows(stream, 32, generator)
        else:                                               # la mitad del lote de cada flujo
            (x1, y1), (x2, y2) = windows(stream, 16, generator), windows(mix_with, 16, generator)
            x, y = torch.cat([x1, x2]), torch.cat([y1, y2])
        loss = nn.functional.cross_entropy(model(x).reshape(-1, len(vocab)), y.reshape(-1))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return loss.item()

base = TinyGPT(len(vocab))
print(f"Parámetros: {sum(p.numel() for p in base.parameters()):,}")
print(f"Pérdida tras el preentrenamiento: {train(base, pretraining_stream, 3000, 3e-3):.3f}")

Parámetros: 113,280

Pérdida tras el preentrenamiento: 0.407

113.280 parámetros, unos miles más que en el cuaderno GPT porque el vocabulario es mayor, 142 tokens. Tras 3.000 pasos, la pérdida es 0,407; como allí, no puede llegar a cero, porque qué diagnóstico, qué consejos y qué nota vienen a continuación es aleatorio.

Paso 4: Lo que sabe el modelo preentrenado, y lo que no hace

---

Un modelo preentrenado solo sabe completar texto. Le damos el comienzo de una nota, "{diagnóstico} se trata con", para cada uno de los 16 diagnósticos, y el comienzo de una hoja hasta "tome". Después le hacemos una pregunta con el formato del paso 5 y vemos cómo continúa.

def continue_text(model, words, max_new=10):
    words = ['<s>'] + list(words)
    new = []
    model.eval()
    with torch.no_grad():
        while len(new) < max_new:
            token = vocab[model(torch.tensor([[token_id[w] for w in words + new]]))[0, -1].argmax()]
            if token in ('.', '</s>'):
                break
            new.append(token)
    model.train()
    return " ".join(new)

def sheet_drug_accuracy(model, n=200):
    """Proporción de hojas de alta nuevas en las que el modelo escribe el fármaco correcto tras 'tome'."""
    sheet_rng, right = np.random.default_rng(777), 0
    for _ in range(n):
        words, diagnosis = discharge_sheet(sheet_rng)
        right += continue_text(model, words[:words.index(TAKE) + 1]).startswith(cases[diagnosis][0])
    return right / n

note_start = note_templates[0].split('{drug}')[0]
completed = sum(continue_text(base, note_start.format(dx=d).split()) == drug for d, (drug, _) in cases.items())
print(f"Completa '{note_start.format(dx='...')}' con el fármaco correcto: {completed} de {len(cases)} diagnósticos")
print(f"Escribe el fármaco correcto en hojas nuevas: {sheet_drug_accuracy(base):.1%}")
question, answer = question_answer('reflujo', drug_questions[0])
prompt = question[:question.index(ANSWER[0]) + 2]
print(f"\n{' '.join(prompt)} ...  ->  {continue_text(base, prompt)!r}   (la respuesta sería {answer!r})")

Completa '... se trata con ' con el fármaco correcto: 16 de 16 diagnósticos

Escribe el fármaco correcto en hojas nuevas: 100.0%

pregunta : ¿ qué debe tomar un paciente con reflujo ? respuesta : ...  ->  'insuficiencia cardiaca'   (la respuesta sería 'omeprazol')

El modelo preentrenado conoce los hechos: completa la nota con el fármaco correcto para los 16 diagnósticos y escribe el fármaco correcto en el 100,0% de las hojas nuevas. Pero si se le hace una pregunta, no la responde: tras "respuesta :" escribe "insuficiencia cardiaca", un diagnóstico en lugar de un fármaco, porque nunca ha visto una pregunta seguida de su respuesta. Es lo que dice el libro de un modelo preentrenado: no está optimizado para ninguna tarea más allá de completar texto.

Paso 5: Los datos de fine-tuning: preguntas sobre 10 de los 16 diagnósticos

---

Cinco plantillas de preguntas, tres sobre el fármaco y dos sobre el signo de alarma, para los 10 primeros diagnósticos: 50 pares de pregunta y respuesta. Es un conjunto diminuto comparado con los 8.000 documentos de preentrenamiento, como dice el libro. Los otros 6 diagnósticos nunca aparecen en una pregunta durante el fine-tuning: las preguntas sobre ellos nos dirán si el modelo sabe usar lo que aprendió en el preentrenamiento.

names = list(cases)
seen, unseen = names[:10], names[10:]
templates = drug_questions + warning_questions
fine_tuning_examples = [question_answer(d, t)[0] for d in seen for t in templates]
fine_tuning_stream = to_stream(fine_tuning_examples)

def question_accuracy(model, diagnoses):
    right = 0
    for d in diagnoses:
        for t in templates:
            question, answer = question_answer(d, t)
            right += continue_text(model, question[:question.index(ANSWER[0]) + 2]) == answer
    return right, len(diagnoses) * len(templates)

print(" ".join(fine_tuning_examples[0]))
print(" ".join(fine_tuning_examples[3]))
print(f"\nEjemplos de fine-tuning: {len(fine_tuning_examples)}, sobre: {', '.join(seen)}")
print(f"Diagnósticos reservados para la prueba: {', '.join(unseen)}")

pregunta : ¿ qué debe tomar un paciente con neumonía ? respuesta : el antibiótico . </s>
pregunta : ¿ cuándo debe volver un paciente con neumonía ? respuesta : dificultad para respirar . </s>

Ejemplos de fine-tuning: 50, sobre: neumonía, migraña, hipertensión, gastroenteritis, cistitis, asma, gota, insuficiencia cardiaca, hipotiroidismo, diabetes
Diagnósticos reservados para la prueba: reflujo, anemia, herpes zóster, trombosis, conjuntivitis, otitis

Paso 6: El fine-tuning, con y sin texto de preentrenamiento

---

Partimos de una copia del modelo preentrenado y la entrenamos con los 50 ejemplos durante 600 pasos, con una tasa de aprendizaje diez veces menor que en el preentrenamiento, para no alejar demasiado los pesos. Lo probamos de dos formas: con lotes formados solo por preguntas, y con la mitad de cada lote tomada del texto de preentrenamiento, una forma habitual de evitar que el modelo olvide lo que sabía (replay). Para cada modelo contamos las respuestas correctas sobre los 10 diagnósticos de los datos de fine-tuning y sobre los 6 nuevos, y comprobamos si sigue escribiendo el fármaco correcto en las hojas de alta.

def report(name, model):
    right_seen, total_seen = question_accuracy(model, seen)
    right_new, total_new = question_accuracy(model, unseen)
    print(f"{name:<44} diagnósticos del fine-tuning {right_seen:>2}/{total_seen} | diagnósticos nuevos {right_new:>2}/{total_new} | "
          f"fármaco correcto en hojas {sheet_drug_accuracy(model):6.1%}")

only_questions = copy.deepcopy(base)
train(only_questions, fine_tuning_stream, 600, 3e-4)
with_replay = copy.deepcopy(base)
train(with_replay, fine_tuning_stream, 600, 3e-4, mix_with=pretraining_stream)

report('Preentrenado, sin fine-tuning', base)
report('Fine-tuning solo con preguntas', only_questions)
report('Fine-tuning con preguntas + texto de preentr.', with_replay)
for d in unseen[:3]:
    question, answer = question_answer(d, drug_questions[0])
    prompt = question[:question.index(ANSWER[0]) + 2]
    print(f"   {' '.join(prompt)}  ->  solo preguntas: {continue_text(only_questions, prompt)!r}, "
          f"con replay: {continue_text(with_replay, prompt)!r}  (correcta: {answer!r})")

Preentrenado, sin fine-tuning                diagnósticos del fine-tuning  0/50 | diagnósticos nuevos  0/30 | fármaco correcto en hojas 100.0%

Fine-tuning solo con preguntas               diagnósticos del fine-tuning 49/50 | diagnósticos nuevos 13/30 | fármaco correcto en hojas  86.5%

Fine-tuning con preguntas + texto de preentr. diagnósticos del fine-tuning 44/50 | diagnósticos nuevos 21/30 | fármaco correcto en hojas 100.0%
   pregunta : ¿ qué debe tomar un paciente con reflujo ? respuesta :  ->  solo preguntas: 'furosemida', con replay: 'omeprazol'  (correcta: 'omeprazol')
   pregunta : ¿ qué debe tomar un paciente con anemia ? respuesta :  ->  solo preguntas: 'las pastillas de hierro', con replay: 'las pastillas de hierro'  (correcta: 'las pastillas de hierro')
   pregunta : ¿ qué debe tomar un paciente con herpes zóster ? respuesta :  ->  solo preguntas: 'aciclovir', con replay: 'aciclovir'  (correcta: 'aciclovir')

Los dos modelos ajustados responden casi todas las preguntas de sus datos. La diferencia está en las 30 preguntas sobre los 6 diagnósticos nuevos. Solo con preguntas, el modelo acierta 13 y empieza a olvidar cómo se escriben las hojas: el fármaco correcto en el 86,5% de ellas. Para el reflujo responde furosemida, el fármaco de uno de sus diagnósticos del fine-tuning. Con la mitad de cada lote tomada del texto de preentrenamiento, acierta 21 de las 30 y sigue escribiendo bien las hojas, el 100,0%, a cambio de algo menos de acierto en sus propias preguntas, 44 de 50. Sus respuestas sobre el reflujo, la anemia y el herpes zóster nombran los fármacos correctos, que solo puede haber aprendido en el preentrenamiento.

Paso 7: Las mismas preguntas, entrenando desde cero

---

¿Es el preentrenamiento lo que marca la diferencia? Entrenamos un modelo con la misma arquitectura, partiendo de pesos aleatorios, con las mismas 50 preguntas, 600 pasos con la tasa de aprendizaje del preentrenamiento.

torch.manual_seed(0)
from_scratch = TinyGPT(len(vocab))
train(from_scratch, fine_tuning_stream, 600, 1e-3)
report('Entrenado desde cero con las preguntas', from_scratch)
question, answer = question_answer(unseen[0], drug_questions[0])
prompt = question[:question.index(ANSWER[0]) + 2]
print(f"   {' '.join(prompt)}  ->  {continue_text(from_scratch, prompt)!r}  (correcta: {answer!r})")

Entrenado desde cero con las preguntas       diagnósticos del fine-tuning 49/50 | diagnósticos nuevos  0/30 | fármaco correcto en hojas   0.0%
   pregunta : ¿ qué debe tomar un paciente con reflujo ? respuesta :  ->  'metformina'  (correcta: 'omeprazol')

El modelo entrenado desde cero aprende las preguntas, 49 de 50, pero no acierta ninguna de las 30 sobre los diagnósticos nuevos: nada en sus datos dice con qué se trata el reflujo, así que responde metformina, uno de los fármacos que ha visto en las preguntas. Y, por supuesto, no sabe escribir hojas de alta. La diferencia entre 0 y 21 de 30 es lo que aporta el preentrenamiento.

Conclusión

---

• El preentrenamiento con texto corriente, prediciendo el siguiente token, mete conocimiento en el modelo, pero el modelo preentrenado solo completa texto: no responde preguntas.
• El fine-tuning con unos pocos ejemplos le enseña un comportamiento nuevo, aquí responder preguntas, con una pequeña fracción del coste del preentrenamiento.
• Lo que el modelo sabe hacer tras el fine-tuning depende de las dos fases, como advierte el libro: el modelo ajustado responde sobre diagnósticos que nunca estuvieron en una pregunta porque aprendió sobre ellos en el preentrenamiento. Un modelo entrenado solo con las preguntas no puede saberlos.
• El fine-tuning también puede hacer que el modelo olvide lo que aprendió. Mezclar texto de preentrenamiento en los lotes de fine-tuning es una forma sencilla de conservarlo.
• El tercer tipo de fine-tuning del libro, la alineación con las preferencias humanas, se construye en el cuaderno Generación de texto y aprendizaje de preferencias del capítulo 8.