Un embedding es una tabla de consulta, y todo lo demás es cómo se rellena

La parte 1 terminó en un índice. El tokenizador convierte el texto en ids, cada id es un número de fila y esa fila está en una matriz con la que nace el modelo. Ahora toca leer esa matriz, que es menos de lo que parece, apenas una tabla de consulta indexada por enteros y con contenido aprendido. La ciencia está en cómo llega ahí ese contenido, y la ingeniería, en lo que ya no puedes hacer con la tabla una vez rellena.

Para esta entrada he leído dos ficheros. El primero es el forward de torch.nn.Embedding, que se reduce a una sola llamada, y el segundo, el módulo de pooling de sentence-transformers, donde se fabrica la otra clase de embedding, la que guardas en un índice. Después he montado un índice de 48 líneas y lo he roto adrede de tres maneras para medir qué falla.

Una matriz y un índice

Si alguna vez has programado una tabla hash, llevas escrita casi toda una capa de embedding. nn.Embedding(V, d) reserva una matriz de V filas y d columnas, y consultar el id i devuelve la fila i. En eso consiste toda la pasada hacia delante, como muestra el propio código de torch en sparse.py:188-197.

def forward(self, input: Tensor) -> Tensor:
    return F.embedding(
        input,
        self.weight,
        self.padding_idx,
        self.max_norm,
        self.norm_type,
        self.scale_grad_by_freq,
        self.sparse,
    )

F.embedding es un gather, y no hace falta fiarse porque se puede comprobar. Con un Embedding(50257, 768) recién creado y tres ids, torch.equal(emb(ids), emb.weight[ids]) da True. No hay transformación, ni hash, ni similitud de por medio. A cada id le corresponde su fila.

Donde deja de parecerse a una tabla hash es en el tamaño. La tabla de entrada de GPT-2, wte.weight, tiene 50 257 filas por 768 columnas, o sea, 38,6 millones de números que ocupan 154 MB en fp32, y cada uno de ellos es un parámetro entrenado. all-MiniLM-L6-v2, el modelo de frases que uso en el resto de la entrada, resuelve lo mismo con una tabla de 30 522 por 384. Si la fila 11 241 de la tabla de GPT-2 es el vector de " token", es porque ahí lo manda el tokenizador, como se vio en la parte 1. El contenido de la fila se aprendió, pero la dirección ya estaba fijada antes de que empezara el entrenamiento.

Al ingeniero que viene del software le sorprende que las claves no signifiquen nada. La clave de una tabla hash significa algo por sí misma, mientras que la fila 11 241 solo tiene sentido dentro del modelo que la entrenó. Al científico de datos le sorprende justo lo contrario. Lo que con word2vec era un artefacto que se produce y luego se usa, aquí es un parámetro del propio modelo y se actualiza en el mismo paso de gradiente que todo lo que tiene encima. Nunca sale de ahí.

Dos maneras de rellenar la tabla

Hay dos maneras de darle números a una fila, y de cuál se use depende lo que puedas hacer con el resultado.

La primera es la de antes, con la tabla como parámetro de un modelo de lenguaje. Las filas empiezan con valores aleatorios y les van llegando los gradientes a través de las capas de encima. Al acabar el entrenamiento, cada fila es lo que haya ayudado al modelo a predecir el siguiente token. Nadie diseña ese espacio, que surge como efecto secundario. Del gradiente que lo rellena se ocupa la parte 8.

La segunda es la que despliega en la práctica la mayoría de los ingenieros. Coges un modelo entrenado, le pasas un texto y te quedas con la salida como vector de ese texto. Esa fila no se guarda en ningún sitio. Se calcula cada vez que hace falta y depende de la entrada entera, de todos sus ids a la vez.

La palabra bank tiene una fila en la tabla de entrada de ese modelo de frases, la 2924. Se le suma una tabla de posiciones, pero la fila de la palabra es la misma en «she sat on the river bank» y en «she opened a bank account». A la salida, en cambio, los dos vectores de ese mismo token tienen un coseno de 0,77 entre sí. Se consulta la misma fila y sale un resultado distinto, porque entre la tabla y la salida hay una posición y seis capas. Eso es lo que quiere decir «contextual», y ahí está toda la diferencia entre un embedding estático y el que metes en un índice de búsqueda.

Lo que sale es un vector por token, pero una frase necesita uno solo, así que algo tiene que condensar en uno los T vectores de los tokens. De eso se encarga el pooling. En sentence-transformers todas las variantes están en un mismo sitio. En la ruta por defecto, con padding, la rama por la que pasa un modelo con pooling de media es pooling.py:200-213.

elif mode in ("mean", "mean_sqrt_len_tokens"):
    if mean_sum is None:
        mask = attention_mask.unsqueeze(-1).expand_as(token_embeddings).to(token_embeddings.dtype)
        mean_sum = (token_embeddings * mask).sum(dim=1)
        # ... mean_mask is the count of real tokens (or a supplied per-token weight sum), clamped away from zero
    if mode == "mean":
        output_vectors.append(mean_sum / mean_mask)

Suma los tokens que la máscara de atención marca como reales y divide entre cuántos había. En esa función conviven seis modos, cada uno una manera distinta de reducir los mismos T vectores. Por ejemplo, cls se queda con el primero, lasttoken con el último y max con el mayor de cada dimensión. Cada modelo se entrena con uno de ellos, y esa elección forma parte del artefacto.

Lo mismo vale para el paso siguiente. El modules.json del propio modelo pequeño enumera lo que se ejecuta al llamar a encode.

[
  {"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"},
  {"idx": 1, "name": "1", "path": "1_Pooling", "type": "sentence_transformers.models.Pooling"},
  {"idx": 2, "name": "2", "path": "2_Normalize", "type": "sentence_transformers.models.Normalize"}
]

Primero va el Transformer, luego el pooling y al final la normalización a longitud unitaria. En el tercero no hay acuerdo ni siquiera dentro de la familia: paraphrase-MiniLM-L6-v2 viene sin él, y sus vectores en bruto tienen normas de alrededor de seis. De los tres módulos, los dos últimos son los que un cliente suele reimplementar, y si lo hace de otra manera, la respuesta cambia. Cuánto cambia se mide dos secciones más abajo.

La cercanía es un producto escalar

Cuando todos los vectores tienen longitud unitaria, la similitud coseno es un producto escalar, y buscar el vecino más cercano se reduce a multiplicar matrices y ordenar. GPT-2 lee su propia tabla del mismo modo a la salida, cuando calcula los logits del siguiente token como el producto escalar con cada fila de wte. La parte de recuperación no tiene más, y por eso la factura del índice es pura aritmética. Diez millones de documentos con 384 dimensiones ocupan 15,4 GB en fp32 sin contar la estructura del índice; con 768 son 30,7 GB, y con 1536, la dimensión de un conocido modelo de embeddings que se sirve por API, 61,4 GB. La dimensión es un coste que pagas tú por cada documento mientras exista el índice, pero que fijó quien entrenó el modelo.

El número que devuelve la búsqueda sirve para menos de lo que parece. En mi corpus de 48 líneas, la mejor coincidencia para «how do I undo a bad release» tiene un coseno de 0,30 con la consulta y es la línea correcta, el procedimiento de rollback. La mejor para «the migration is stuck waiting on a lock» saca 0,71 y también acierta. Un coseno es una posición en un espacio cuya escala ha elegido el modelo, y no sirve como medida de confianza. Compara el orden de los resultados dentro de un mismo modelo y nunca compares el número en bruto entre dos.

No hay migración de esquema

Lo que viene ahora no tiene equivalente en el software. Cuando cambias el esquema de una base de datos, escribes una migración, la ejecutas y las filas viejas quedan con la forma nueva. Un espacio de embeddings no tiene migración. Dos modelos producen dos espacios y ninguno trae un mapa hacia el otro, porque nadie los diseñó y cada uno es donde fue a parar el entrenamiento. Puedes ajustar un mapa a posteriori, y hay quien lo hace. El resultado es una aproximación aprendida, con su propio error, y eso no es una migración. Con un mapa lineal ajustado sobre cuatro mil frases ajenas al corpus, el índice desfasado que mido más abajo pasa del 60 % al 70 % y no sube más. Si recalculas los embeddings con un modelo nuevo, cada vector que tienes guardado queda en un sistema de coordenadas que no usa nadie más.

Esperaba que eso se tradujera en resultados basura, y lo que sale es algo peor.

La prueba usa 48 líneas de un runbook, ocho de cada tema: despliegues, autenticación, facturación, bases de datos, observabilidad y guardias. Las he indexado con el modelo L6 y he lanzado las consultas con all-MiniLM-L12-v2, la actualización a la que recurriría un ingeniero (misma familia y mismas 384 dimensiones, así que el índice acepta los vectores sin rechistar).

Frente a lo que devolvería el índice actualizado, el desfasado conserva el 60 % de los cinco primeros resultados. Al azar se conservaría el 10 %, y adivinando cinco líneas del tema correcto, el 35 %. El índice viejo consultado con el modelo viejo conserva el 77 %: el desajuste cuesta diecisiete puntos y los otros veintitrés se deben a que los dos modelos no coinciden en la respuesta. Con «user keeps getting logged out», el primer resultado sigue siendo la línea de la caducidad de sesión, pero el tercero es la política de retención de las copias de seguridad de la base de datos. Nada en la respuesta te dice cuál es cuál.

Para ser justos con la cifra, un 60 % no es un índice revuelto, y para explicarlo no basta con que los dos modelos trabajen en inglés: la razón está en el linaje. Ambos se ajustaron a partir del mismo checkpoint padre, del que el L6 conserva solo una capa de cada dos, con lo que sus espacios empezaron en las mismas coordenadas. La misma línea codificada con uno y otro tiene, de media, un coseno de 0,54 consigo misma, muy por encima de lo que darían dos vectores ortogonales. Dos modelos que solo compartieran los datos de entrenamiento no compartirían coordenadas, y un índice desfasado entre ellos debería quedar cerca del azar. Además, seis temas bien delimitados de ocho líneas cada uno son el corpus más favorable que podría pedir un índice desfasado. No he medido un índice en producción, y esperaría que conservara menos. Lo que sí respalda la cifra es la forma del fallo, que es parcial, verosímil y silencioso.

El módulo de pooling produce el mismo efecto sin cambiar de modelo. Con los mismos pesos, si cambias mean por cls al indexar y al consultar, los cinco primeros conservan el 50 % de la respuesta con pooling de media. Si los mezclas, con media en el índice y cls en la ruta de consulta, conservan el 70 %. Sale mejor que con cls en los dos lados porque uno de ellos sigue usando el pooling con el que se entrenó el modelo. Un cliente que reimplementa el pooling ha cambiado el modelo.

Avanza la reindexación y luego cambia el modelo que codifica las consultas. Con el modelo nuevo, el recall sube a medida que migra la tabla. Con el viejo empieza en 0,77 porque dos modelos usados de forma coherente ya discrepan en una cuarta parte de las respuestas. Después cae hasta 0,63 en la parte central de la migración, no vuelve al 0,77 hasta las últimas líneas y termina por debajo, en 0,73. Para entonces, cada línea del índice la ha escrito un modelo que la ruta de consulta no usa.

Esa segunda curva es la que te encuentras en producción, donde el índice sale de un proceso por lotes, el codificador de consultas es un servicio desplegado y cada uno se actualiza un día distinto.

Qué vigila el ingeniero de ML

Son cuatro hábitos, todos baratos al lado de una reindexación.

El modelo de embeddings forma parte del índice. Guarda en los metadatos del índice el hash del modelo, el modo de pooling, el paso de normalización y la longitud máxima de secuencia, y rechaza cualquier consulta cuyo codificador no coincida. El modules.json de antes es la mitad de ese contrato; el modo de pooling está a su lado, en 1_Pooling/config.json, y el hash lo añades tú. La longitud máxima importa. El modelo L6 lee 256 tokens y la versión L12 solo 128, de modo que el cambio que rompió el índice también redujo a la mitad, sin avisar, lo que ve cada vector. Todo eso tiene que ir con el índice, fuera de la configuración de cualquier cliente.

Corte limpio, sin reindexación gradual. Construye el índice nuevo junto al viejo, cambia consultas e índice a la vez, detrás de un único flag, y borra el viejo al terminar. Una reindexación gradual con la ruta de consulta en marcha es justo la curva que cae en la demo.

Mide el solapamiento de vecinos y olvídate del coseno. Mantén un conjunto fijo de consultas con su top-k esperado y calcula con él el recall@k cada vez que cambies de modelo. Un umbral de coseno ajustado con un modelo no significa nada con el siguiente, y el conjunto de vecinos es lo único observable que sobrevive al cambio. Eso sí, el solapamiento mide la deriva y deja fuera la relevancia: con el índice ya migrado, los cinco primeros de la consulta del cierre de sesión todavía incluyen una línea sobre almacenamiento en frío.

La dimensión es una factura. Son N por d por cuatro bytes, sin contar la estructura del índice. La media precisión reduce la factura a la mitad, y algunos modelos recientes te dejan quedarte solo con las k primeras dimensiones. Las dos cosas se pueden aplicar a vectores ya guardados sin recalcular los embeddings, aunque los vectores truncados hay que renormalizarlos y las consultas hay que cortarlas a la misma k. Ambas decisiones van en los metadatos del índice, con todo lo demás. Calcula lo que te costará la dimensión de un modelo mayor con el corpus que tendrás dentro de dos años, porque el de hoy no sirve para esa cuenta.

Esto tampoco es trabajo de modelado. Un índice de embeddings parece una base de datos con una columna de floats hasta el día en que alguien actualiza el modelo que escribió esa columna.