Resumen Técnico: IG-Lens

Planteamiento del Problema

Los transformadores de solo decodificador (decoder-only) se conceptualizan a menudo como procesos de inferencia iterativos donde el flujo residual refina una estimación continua del siguiente token. Surge una pregunta crítica de interpretabilidad: ¿en qué capa específica se decide realmente la probabilidad de un token predicho?

Las herramientas existentes para la lectura por capas (layer-wise readout) no logran responder a esto con precisión en el espacio de probabilidad:

• Logit Lens y Tuned Lens: Informan el nivel de probabilidad en cada capa, pero no proporcionan una descomposición aditiva. Sus estimaciones son sesgadas y no monótonas a través de la profundidad, lo que significa que la suma de las probabilidades por capa no tiene un significado matemático.
• Atribución Directa de Logits (DLA) y Métodos de Embedding: Son aditivos pero operan en el espacio de los logits. Debido a que la función softmax es altamente no lineal, una descomposición aditiva de los logits no se traduce en una descomposición aditiva de las probabilidades, que es la magnitud de interés práctico.
• Conductancia de Capa e Integrated Gradients de Capa (IG): Estos métodos integran gradientes por capa, pero atribuyen cada capa a su propio baseline independiente. Por consiguiente, los términos por capa no se "telescopan" (no se suman) para alcanzar el cambio total en la probabilidad de predicción.

Metodología: IG-Lens

Los autores presentan IG-Lens, un método que proporciona una descomposición aditiva exacta de la probabilidad a través de las capas del transformador.

Concepto Central

IG-Lens aplica una aplicación telescópica de Integrated Gradients (IG) a lo largo de un único camino a través de los estados ocultos desde un baseline hasta la capa final.

1. Definición del Camino: El método selecciona una secuencia de capas $0 = L_0 < L_1 < \dots < L_k = n$ y define un camino de estados ocultos $h_{L_0} \to h_{L_1} \to \dots \to h_{final}$.
2. Lectura de Cabezal Único (Head-Once Readout): Crucialmente, la función de lectura $f(h)$ se define como la aplicación final del unembedding a un único vector oculto: $f(h) = \text{softmax}(\text{head}(\text{norm}(h)))_{y_t}$. Esta función no contiene atención ni MLP; se asume que toda la mezcla entre tokens ya ha ocurrido en el paso hacia adelante (forward pass) que produjo los estados ocultos desvinculados.
3. Atribución Telescópica: La contribución de la capa $L_j$ se define como el segmento IG de $f$ a lo largo de la línea recta desde $h_{L_{j-1}}$ hasta $h_{L_j}$. Por el teorema del gradiente, la integral del gradiente sobre un segmento es igual a la diferencia de los valores de los extremos. Sumar estos segmentos produce una suma telescópica:
   $$\sum_{j=1}^k IGL_j = f(h_{final}) - f(h_{base}) = \Delta p$$
   Esto garantiza que la suma de las atribuciones es exactamente el cambio en la probabilidad objetivo, con el softmax y el LayerNorm integrados en lugar de ser linealizados o congelados.

Estimador Consciente de la Predicción (Prediction-Aware Estimator)

Aunque la integral es exacta, los estimadores de pasos finitos suelen introducir errores de discretización. IG-Lens adopta un estimador consciente de la predicción (inspirado en IDGI) como su predeterminado:

• En lugar de acreditar los pasos basándose en gradientes brutos, acredita cada paso por su cambio observado en la probabilidad objetivo ($f(h^{(s)}) - f(h^{(s-1)})$).
• Para una lectura de probabilidad escalar, esto colapsa cada segmento en una suma telescópica de los valores de los extremos.
• Resultado: La completitud se mantiene hasta la precisión de punto flotante en cualquier número de pasos ($m \ge 1$). Los pasos donde el modelo muestra sensibilidad al gradiente pero no hay cambio en la salida son filtrados naturalmente, eliminando el error de discretización de Riemann.

Implementación

El método utiliza un algoritmo de paso único por lotes (single-pass batched algorithm) que computa el mapa completo de token por capa sin requerir un paso hacia atrás (backward pass/autograd).

• Al apilar todos los puntos de interpolación en un solo tensor y evaluar la función puntual $f$ en un solo paso hacia adelante, el método evita la sobrecarga del cálculo de gradientes.
• Esta eficiencia es posible específicamente debido al diseño de "cabezal único", que asegura la independencia entre los puntos de los tokens en el lote.

Contribuciones Clave

1. Descomposición Aditiva Exacta: IG-Lens es el primer método que proporciona una descomposición aditiva exacta de la probabilidad (no de los logits) a través de las capas, donde la suma de las atribuciones es igual a $\Delta p$.
2. Identidad Telescópica: El artículo proporciona una prueba formal (Proposición 1) de que las integrales de segmento se televascan hacia el cambio total en la predicción, manteniendo el softmax dentro de la trayectoria de integración.
3. Estimación Consciente de la Predicción: El estimador predeterminado asegura la completitud con precisión de máquina independientemente del número de pasos y filtra los pasos de "sensibilidad sin respuesta".
4. Implementación Eficiente: Una evaluación de paso único (forward evaluation) computa el mapa de atribución completo sin llamadas a backward, lo que lo hace computacionalmente eficiente.
5. Posicionamiento: Los autores delinean claramente por qué los métodos existentes fallan en el espacio de probabilidad (aditividad en el espacio de logits, sesgo o falta de telescopía) y posicionan a IG-Lens como la contraparte en el espacio de probabilidad de los métodos en el espacio de logits.

Resultados y Plan de Evaluación

• Verificación Numérica: En pruebas controladas y en un modelo real (Llama-3.2-1B-Instruct), el estimador predeterminado reproduce las diferencias de los extremos de cada segmento hasta $\sim 10^{-7}$ (precisión float32). La suma telescópica se mantiene exacta según el punto flotante, independientemente del número de pasos $m$.
• Salida Ilustrativa: En experimentos con el prompt "What is the capital of Vietnam?", el método logra concentrar la masa de atribución en los segmentos de capa específicos donde el valor del logit-lens transiciona de bajo a alto. La suma de las atribuciones para cada token coincide con $\Delta p$ con precisión de punto flotante.
• Estrategia de Evaluación: Los autores reconocen que el "inicio" (onset) no tiene una verdad de terreno (ground truth) intrínseca. Por lo tanto, proponen validar el método mediante ablación causal: si IG-Lens identifica la capa $L$ como el inicio, la ablación de $L$ debería deprimir la probabilidad objetivo más que la ablación de capas con puntuaciones de inicio más bajas. Esta prueba está diseñada para ser económica de ejecutar y capaz de falsificar la afirmación.

Significado y Reivindicaciones

El artículo posiciona a IG-Lens como una refinación cuantitativa de la visión de inferencia iterativa. A diferencia del Logit Lens, que pregunta "¿cuál es la predicción en ejecución en esta capa?", IG-Lens pregunta "¿cuánta de la probabilidad final es responsabilidad de cada capa?" y garantiza que estas responsabilidades sumen el todo.

Reivindicaciones Modestas y Limitaciones:

• Marginal Condicional: Los autores declaran explícitamente que IG-Lens mide un marginal condicional: la probabilidad adicional que el norm+head puede leer de $h_L$ dado $h_{L-1}$. Este método descarta el efecto de $h_L$ procesado a través de los bloques superiores de atención y MLP.
• Compromiso (Trade-off): El método elige "suma = $\Delta p$" sobre el "efecto causal total a través de todo el paso hacia adelante". Para efectos causales totales, la descomposición de componentes (como DLA) sigue siendo la herramienta apropiada.
• No Monotonicidad: El método maneja trayectorias de probabilidad no monótonas (donde la probabilidad se construye y se deshace parcialmente) utilizando atribuciones con signo, ofreciendo una visión más fiel de las contribuciones de las capas que los métodos basados en valores absolutos.

La cuestión abierta central sigue siendo empírica: si el inicio de IG-Lens predice la capa en la que la ablación causal deprime más la predicción, particularmente para tokens con trayectorias no monótonas.