El logit lens es una técnica de interpretabilidad en aprendizaje automático utilizada para inspeccionar las representaciones internas de los modelos de Transformer (architecture) basados en grandes modelos de lenguaje. Funciona tomando el estado oculto en una capa dada, aplicando la matriz de desincrustación final del modelo (que mapea los estados ocultos a logits sobre el vocabulario) y luego aplicando un softmax para obtener una distribución de probabilidad. Esta distribución revela qué predeciría el modelo si se detuviera en esa capa, ofreciendo una ventana a cómo evolucionan las predicciones capa por capa. El método fue introducido por nostalgebraist en 2020 y desde entonces se ha convertido en una herramienta estándar para estudiar los cálculos internos de modelos como GPT-2 y GPT-3.
A diferencia de los clasificadores de sondeo que requieren entrenar modelos auxiliares, el logit lens utiliza la propia cabeza de salida del modelo, lo que lo convierte en un enfoque de cero parámetros y sin entrenamiento. Es particularmente efectivo para modelos con flujos residuales, ya que los estados ocultos en cada capa son aditivos y pueden proyectarse directamente. La técnica se ha utilizado para identificar cuándo los modelos se comprometen con predicciones específicas, detectar decodificación "tardía" versus "temprana" y analizar fenómenos como el recuerdo factual y el razonamiento aritmético. Sin embargo, su aplicabilidad se limita a modelos con incrustaciones atadas o compartidas, y puede no funcionar bien para modelos con normalización de capas compleja o transformaciones no lineales.
Mecanismo e Implementación
El logit lens opera sobre el flujo residual de un transformer. En un transformer típico, cada capa añade su salida al flujo residual, por lo que el estado oculto en la capa \( l \) es la suma de la incrustación inicial y todas las salidas de capas hasta \( l \). Los logits finales se calculan aplicando una normalización de capas (si está presente) y luego la matriz de desincrustación \( W_U \). El logit lens aplica la misma desincrustación a los estados ocultos intermedios, a menudo después de aplicar la normalización de capas final (o una transformación afín aprendida) para tener en cuenta el cambio de distribución.
Para modelos con incrustaciones atadas (donde la incrustación de entrada y la desincrustación de salida comparten pesos), la proyección es sencilla. Para modelos con matrices de desincrustación separadas, la técnica sigue funcionando siempre que la desincrustación esté disponible. El resultado es una secuencia de distribuciones de probabilidad sobre el vocabulario para cada capa, que puede visualizarse como un mapa de calor o usarse para calcular métricas como la "diferencia de logits" entre tokens candidatos.
Aplicaciones en Interpretabilidad
El logit lens se ha aplicado a una variedad de tareas de interpretabilidad. Un uso común es rastrear el desarrollo de predicciones a través de las capas. Por ejemplo, en un modelo al que se le pide completar "La capital de Francia es __", el logit lens podría mostrar que las capas tempranas asignan alta probabilidad a "París" solo después de varias capas, revelando cuándo el modelo recupera el hecho relevante. Los investigadores han utilizado esto para identificar "cabezas de inducción" y otros patrones de atención que contribuyen a comportamientos específicos.
Otra aplicación es detectar decodificación "tardía", donde un modelo puede predecir inicialmente un token pero luego cambiar de opinión. Al examinar el logit lens en cada capa, se puede ver cuándo la predicción final está "bloqueada" y cuándo aún es flexible. Esto tiene implicaciones para comprender la confianza del modelo y para diseñar intervenciones.
La técnica también se ha utilizado para estudiar la aritmética y el razonamiento. En modelos como GPT-3, el logit lens puede mostrar que las capas intermedias representan sumas o acarreos intermedios, proporcionando evidencia de cálculo en múltiples pasos. Esto ha llevado a más trabajo en interpretabilidad mecanicista, como el marco de "cabezas de inducción" y "circuitos" de Anthropic y otros grupos de investigación.
Limitaciones y Críticas
A pesar de su utilidad, el logit lens tiene varias limitaciones. Primero, asume que la matriz de desincrustación es una proyección significativa para todas las capas, lo que puede no cumplirse si los estados ocultos del modelo sufren transformaciones significativas (por ejemplo, normalización de capas) que no se tienen en cuenta. Algunos modelos, especialmente aquellos con arquitecturas de pre-normalización, requieren aplicar la normalización de capas final a los estados intermedios, pero esto no siempre es suficiente.
Segundo, el logit lens es menos efectivo para modelos con vocabularios grandes o para tareas que requieren predicciones de múltiples tokens, ya que solo muestra la distribución del siguiente token. También falla en capturar el estado completo del modelo, ya que ignora la contribución de las subcapas de atención y de avance más allá de su efecto aditivo.
Tercero, la técnica puede producir resultados engañosos si el modelo usa una cabeza de salida separada que no está alineada con el flujo residual. En tales casos, los logits proyectados pueden ser ruidosos o poco informativos. Los investigadores han propuesto variantes como el "tuned lens" que aprenden una transformación afín por capa para mejorar la alineación, pero requieren datos de entrenamiento adicionales.
Relación con Otros Métodos de Interpretabilidad
El logit lens es parte de una familia más amplia de técnicas de interpretabilidad que incluyen clasificadores de sondeo, parcheo de activaciones y rastreo causal. Los clasificadores de sondeo entrenan un modelo lineal o no lineal sobre estados ocultos para predecir características específicas, pero requieren supervisión y pueden no reflejar el cálculo real del modelo. En contraste, el logit lens utiliza la propia cabeza de salida del modelo, lo que lo hace más directo.
El parcheo de activaciones implica intervenir en los estados ocultos para ver cómo afectan la salida, lo que puede usarse para validar hallazgos del logit lens. El rastreo causal, popularizado por David Kaplan y otros, utiliza ejecuciones corruptas para identificar qué capas son responsables de hechos específicos, y el logit lens puede complementar esto mostrando el token predicho en cada capa.
Otro método relacionado es la vista de "flujo residual", que trata al transformer como una serie de actualizaciones aditivas. El logit lens es un ajuste natural para esta vista, ya que proyecta el flujo residual en cada punto. Esto ha inspirado herramientas como el "Transformer Debugger" y otras bibliotecas de visualización.
Extensiones y Variantes
Se han propuesto varias extensiones del logit lens. El "tuned lens" aprende una transformación afín por capa para alinear mejor los estados ocultos con la desincrustación, mejorando el rendimiento en modelos donde el logit lens estándar falla. Otra variante es el "softmax lens", que aplica una temperatura a los logits para afilar o aplanar la distribución, facilitando ver predicciones de baja probabilidad.
Para modelos con múltiples cabezas de salida o decodificadores, el logit lens puede aplicarse a cada cabeza por separado. En modelos codificador-decodificador, la técnica puede aplicarse a los estados ocultos del decodificador, pero es menos claro cómo aplicarla al codificador.
El trabajo reciente también ha explorado el uso del logit lens para analizar modelos multimodales, donde el vocabulario incluye tokens de imagen u otras modalidades. Sin embargo, esto sigue siendo un área activa de investigación.
Consideraciones Computacionales
El logit lens es computacionalmente eficiente, ya que solo requiere una pasada hacia adelante y una multiplicación de matrices en cada capa. Para un modelo con \( L \) capas y un tamaño de vocabulario \( V \), el costo adicional es \( O(L \cdot V \cdot d) \), donde \( d \) es la dimensión oculta. Esto es insignificante en comparación con el costo de la pasada hacia adelante en sí.
Sin embargo, almacenar los logits para todas las capas puede ser intensivo en memoria, especialmente para vocabularios grandes. En la práctica, se pueden calcular los top-k tokens en cada capa para reducir el uso de memoria. La técnica está implementada en varias bibliotecas de código abierto, incluida la biblioteca TransformerLens y las herramientas de interpretabilidad de Hugging Face.
Direcciones Futuras
A medida que los modelos de aprendizaje profundo continúan creciendo en escala, herramientas de interpretabilidad como el logit lens se vuelven cada vez más importantes. El trabajo futuro puede centrarse en adaptar la técnica a nuevas arquitecturas, como mezclas de expertos o modelos de espacio de estados, y en integrarla con pipelines de interpretabilidad automatizada. El logit lens también se está utilizando en investigación de seguridad para detectar cuándo los modelos están "pensando" en salidas dañinas, lo que podría informar estrategias de alineación.
En general, el logit lens sigue siendo una herramienta simple pero poderosa para mirar dentro de la caja negra de las redes neuronales, y sus aplicaciones probablemente se expandirán a medida que el campo de la interpretabilidad mecanicista madure.