English · Español
02 — Máscaras de logits: la derivación¶
Restringir el muestreo a un subconjunto \(\mathcal{L}\) de tokens legales es exactamente equivalente a multiplicar los logits por \(-\infty\) fuera de \(\mathcal{L}\) antes del softmax. La distribución resultante es la condicional sobre \(\{t \in \mathcal{L}\}\). No hay magia, solo Bayes.
Esta es la página de teoría que carga el peso. Deriva la fórmula, interiorízala, y el resto de la fase es fontanería.
Planteamiento¶
El modelo produce un vector de logits \(z \in \mathbb{R}^{|V|}\) en cada paso, condicionado por el prefijo \(x_{<i}\). La distribución de muestreo no restringida es
Tenemos una restricción: solo los tokens en un subconjunto \(\mathcal{L}_i \subseteq V\) son legales en el paso \(i\) (donde \(\mathcal{L}_i\) depende del prefijo y del grammar). Queremos muestrear de
Aplica la definición de probabilidad condicional:
Numerador: la probabilidad original para tokens legales, cero para ilegales. Denominador: la masa total sobre tokens legales, usada para renormalizar.
El truco de la máscara de logits¶
Computar \(p_\text{constr}\) con esa fórmula requiere: 1. Softmax de los logits completos (\(|V|\) exp, \(|V|\) suma). 2. Multiplicar por la máscara indicadora. 3. Renormalizar por la suma sobre tokens legales.
Eso son 3 pasadas sobre el vocabulario. El truco de la máscara es empujar el indicador dentro de los logits antes del softmax:
Entonces \(\exp(\tilde{z}_t) = \exp(z_t) \cdot \mathbb{1}[t \in \mathcal{L}_i]\) (ya que \(\exp(-\infty) = 0\)), y el softmax de \(\tilde{z}\) es exactamente \(p_\text{constr}\). Dos pasadas: enmascarar y luego softmax-con-renormalización. Distribución idéntica.
En código:
masked_logits = logits + mask # mask es 0 para legal, -inf para ilegal
probs = softmax(masked_logits)
Ese mask es exactamente el entregable que construye la Fase 30. Todo lo de downstream — temperature, top-k, top-p, el sampler de la Fase 21 — opera sobre masked_logits (o probs) sin cambios.
Consideraciones numéricas¶
-inf es un valor real de numpy.float32 (np.float32(-np.inf)), y exp(-inf) == 0.0 exactamente. No hay problema de redondeo.
El peligro es que todos los logits queden enmascarados. Entonces masked_logits es todo -inf, el softmax computa 0/0 y produce NaN. Este es el caso "\(\mathcal{L}_i\) vacío" (pitfall 3 en PHASE_30_PLAN.md §5). Código defensivo: aserta (mask > -inf).any() antes del softmax; lanza una excepción NoLegalContinuation si el grammar se ha pintado a sí mismo en una esquina.
Para la estabilidad usamos el softmax estándar con el truco de la sustracción del -max (de la Fase 2). Cuidado: el truco del -max sustrae el máximo de los logits enmascarados, que es -inf si todos están enmascarados. La comprobación defensiva de arriba evita que esto llegue al softmax.
Composición con estrategias de muestreo¶
El sampler de la Fase 21 puede hacer temperature, top-k, top-p, penalizaciones de repetición. ¿Cómo se compone la máscara con estas?
Temperature. La temperature escala los logits antes del softmax: \(z'_t = z_t / T\). La máscara es \(\{0, -\infty\}\); escalar cualquiera devuelve \(\{0, -\infty\}\) (\(-\infty / T = -\infty\) para \(T > 0\)). Así que máscara y temperature conmutan. Aplica cualquiera primero. Convención: aplicar la máscara primero (la operación más barata; las entradas enmascaradas se saltan la multiplicación por temperature si lo haces con cabeza).
Top-k. Top-k se queda con los \(k\) logits más altos. Si aplicas top-k primero y luego la máscara, podrías acabar con menos de \(k\) tokens legales (algunos de los top-\(k\) se enmascaran). Si enmascaras primero y luego top-k, eliges el top-\(k\) entre los tokens legales. Esto último es lo que quieres. Enmascarar antes de top-k.
Top-p. Top-p (nucleus) se queda con el conjunto más pequeño de tokens cuya suma de probabilidad ≥ \(p\). Mismo argumento: enmascarar primero, luego top-p sobre la distribución renormalizada. Si no, top-p podría seleccionar un token ilegal que se descarta, dejando el nucleus indefinido.
Penalización de repetición. Multiplica los logits de tokens emitidos recientemente por algún factor. Conmuta con la máscara (ilegal sigue siendo ilegal). El orden no importa, pero aplicar la máscara primero es más barato.
Regla general: la máscara es la operación más exterior; todo lo demás está dentro. En código:
def sample(logits, mask, temperature, top_p, ...):
logits = logits + mask # 1. máscara
logits = logits / temperature # 2. temperature
logits = apply_repetition_penalty(...) # 3. penalización de repetición
probs = softmax(logits) # 4. softmax
probs = apply_top_p(probs, top_p) # 5. top-p
token = sample_from(probs) # 6. muestrear
return token
El lab de la Fase 30 implementa exactamente este orden y añade un test que verifica que un token ilegal nunca es muestreado.
El diagnóstico KL¶
¿Cuánto distorsionó la máscara al modelo? Mide con divergencia KL:
Una forma más simple usando \(Z = \sum_{t \in \mathcal{L}_i} p(t)\) (la masa legal total):
Derivación: \(p_\text{constr}(t) = p(t) / Z\) para \(t \in \mathcal{L}_i\), así que \(\log(p_\text{constr} / p) = -\log Z\). La suma ponderada por \(p_\text{constr}\) es simplemente \(-\log Z\) (es constante).
Interpretación. Si \(Z \approx 1\) (el modelo ya iba a emitir un token legal), \(\mathrm{KL} \approx 0\) — el enmascarado es un no-op. Si \(Z \approx 0\) (el modelo quería emitir algo ilegal), \(\mathrm{KL} \to \infty\) — estamos forzando al modelo muy lejos de su distribución preferida.
Un \(\mathrm{KL}\) consistentemente alto a lo largo de los pasos de decode significa que el modelo está peleando contra el grammar. Esto es una señal, no un fallo: te dice que el modelo no fue entrenado en este formato y se le está coaccionando, lo que puede producir salida semánticamente pobre aunque parsee. Lo registramos en experiments/30-mask-overhead/.
Computar la máscara¶
La máscara depende del estado actual del parser del grammar dado el prefijo emitido hasta el momento. Para nuestro caso de uso de JSON Schema:
state := uno de {
EXPECT_OPEN_BRACE,
EXPECT_KEY_OPEN_QUOTE,
EXPECT_KEY_CHARS,
EXPECT_KEY_CLOSE_QUOTE,
EXPECT_COLON,
EXPECT_VALUE_BY_KEY, # depende de qué clave
EXPECT_COMMA_OR_CLOSE,
DONE,
}
En cada paso, dado state, enumeramos el vocabulario y preguntamos: "si emito el token \(t\), ¿en qué estado acaba el parser?". Si la respuesta es un estado válido (o un camino válido a un estado válido), \(t\) es legal; en otro caso, enmascáralo.
La complicación son los tokens multi-carácter. Un token como "verb": es un único token BPE pero cruza EXPECT_KEY_OPEN_QUOTE → EXPECT_KEY_CHARS → EXPECT_KEY_CLOSE_QUOTE → EXPECT_COLON → EXPECT_VALUE_BY_KEY (cinco transiciones de estado). La lógica de la máscara debe simular la secuencia completa y solo aceptar si cada estado intermedio es legal y el estado final es uno al que queríamos llegar.
Para la Fase 30 implementamos esta simulación de forma ingenua: para cada token candidato, lo decodificamos a un string, ejecutamos el parser carácter a carácter, aceptamos o rechazamos. O(\(|V| \cdot \bar{L}\)) por paso, donde \(\bar{L}\) es la longitud media del token. Lento pero correcto.
Una comprobación de cordura: identidad cuando \(\mathcal{L}_i = V\)¶
Si todos los tokens son legales (\(\mathcal{L}_i = V\)), la máscara es toda ceros y \(p_\text{constr} = p\). Enmascarar es un no-op en este caso. Los tests de la Fase 30 verifican esto: una máscara "permisiva" que permite todo produce muestras idénticas a no usar máscara (módulo diferencias de estado de RNG, que nuestro sampler determinista evita).
Esta es la primera cosa que testear. Una implementación de máscara que falla este test está rota.
Una segunda comprobación de cordura: degenerada cuando \(\mathcal{L}_i = \{t^*\}\)¶
Si solo un token \(t^*\) es legal, \(p_\text{constr}\) es una masa puntual sobre \(t^*\). El muestreo siempre devuelve \(t^*\). El decoder está efectivamente forzado en este paso. El esquema de conjugación tiene varios de estos pasos (p. ej., tras {, solo " es legal; tras "verb", solo : es legal).
Estos pasos "forzados" son donde el modelo no tiene voz. Los pasos interesantes están dentro de los campos valor, donde el modelo elige el verbo, el tense o la person reales del enum cerrado. La máscara allí permite solo los valores legales del enum; el modelo elige entre ellos.
Lo que significa esto operacionalmente¶
Para el esquema de conjugación, grandes tramos de la salida están forzados — la puntuación, los nombres de campo. El modelo solo "crea" dentro de los valores, e incluso ahí la elección es de un enum pequeño (20 verbos, 5 tenses, 3 persons). Esto reduce dramáticamente la entropía de la secuencia generada. En la práctica, la generación estructurada a menudo emite decenas de tokens de sobrecoste (el andamiaje JSON) por cada predicción de contenido de pocos tokens. Eso está bien — los tokens están forzados, así que el coste es solo un matmul por token en tiempo de decode, no trabajo de inferencia significativo.
Una trampa: el modelo no "sabe" que está siendo enmascarado¶
La distribución del siguiente token del modelo se computa asumiendo continuación no restringida. Su estado oculto codifica su expectativa del futuro, no la del grammar. Si el grammar lo fuerza por un camino que consideraba improbable, la KV cache del modelo de pasos anteriores puede ser engañosa para la predicción restante. En la práctica esto está bien para esquemas pequeños; en teoría es una fuente de distribution shift que los sistemas de producción a veces evitan con "enseñar al modelo a emitir el formato nativamente vía ajuste fino (fine-tuning)" (territorio de la Fase 28).
Para la Fase 30 lo ignoramos — nuestro modelo es minúsculo y el esquema es pequeño.
Lo que esta fase NO cubre¶
- Corrección de distribution shift. Un "speculative re-prefill" con un modelo ajustado por ajuste fino (fine-tuning) que ya conoce el formato es una extensión de la Fase 28 (LoRA); no lo hacemos aquí.
- Estado de máscara por beam. La Fase 30 solo hace greedy / top-p. La interacción beam search × máscara es un tema aparte; mencionada en
PHASE_30_PLAN.md§7. - Enmascarado a mitad de token (sub-BPE). La versión de producción de Outlines divide tokens aún más cuando cruzan transiciones de grammar de forma incómoda; nosotros aceptamos el token que nos dé el BPE y validamos en la frontera de token.
- Grammars no-JSON. Nuestra máscara es específica de JSON Schema. Las descripciones de CFG / GBNF están en
03-grammar-as-dfa.md, sin implementar.
Siguiente: theory/03-grammar-as-dfa.md — cómo las implementaciones de producción precompilan el grammar en un autómata, y por qué no lo hacemos aquí.