Resumen Técnico: TN-SHAP-G

Declaración del Problema

Explicar las predicciones de las redes neuronales de grafos (GNN) de caja negra requiere atribuir importancia a los nodos individuales y a sus interacciones. Aunque los valores de Shapley y los índices de interacción de orden superior ofrecen atribuciones axiomáticas y fundamentadas, su computación exacta es intratable para grafos, escalando como $O(2^n)$ donde $n$ es el número de nodos. Los métodos de aproximación existentes enfrentan compromisos significativos: los explicadores específicos de la arquitectura (por ejemplo, GNNExplainer) no se aplican a modelos de caja negra generales, mientras que los enfoques basados en muestreo (por ejemplo, GraphSVX, SHAP-IQ) requieren miles de evaluaciones del modelo y sufren de una alta varianza, particularmente para interacciones de orden superior. Además, estos métodos de muestreo a menudo no logran escalar a grafos más grandes donde la enumeración exacta es imposible.

Metodología: TN-SHAP-G

Los autores proponen TN-SHAP-G, un marco que explota las propiedades estructurales de las entradas con estructura de grafo para computar valores de Shapley e interacciones de manera eficiente y determinista. La idea central es aprender un sustituto de red de tensores (TN) compacto y alineado con el grafo que aproxime el comportamiento de la entrada enmascarada del predictor.

1. El Juego de Grafo Enmascarado

El método formaliza la explicación de grafos como un juego cooperativo inducido por el enmascaramiento de nodos. Para un grafo de entrada fijo $G=(V, E)$ con características $X$ y una línea base $X^0$, una coalición $S \subseteq V$ define una matriz de características enmascarada $X_S$ donde los nodos en $S$ conservan sus características originales y los nodos fuera de $S$ son reemplazados por características de línea base. El valor del juego es $\nu(S) = f(G, X_S)$, donde $f$ es el predictor de caja negra.

2. Extensión Multilineal y Contracción de Tensores

La computación directa de los valores de Shapley requiere sumar sobre $2^n$ coaliciones. Los autores utilizan la extensión multilineal $\hat{\nu}(z)$, una relajación continua del juego discreto definido sobre $z \in [0, 1]^n$.

• Formulación de Tensores: La extensión multilineal puede expresarse exactamente como una contracción de tensores: $\hat{\nu}(z) = \sum_{s \in \{0,1\}^n} T_s \prod_{v=1}^n b_v(z_v)^{s_v}$, donde $T$ es el tensor de coalición y $b_v(z_v)$ son vectores selectores de Bernoulli.
• Recuperación Determinista: Los valores de Shapley y los índices de interacción corresponden a integrales de las derivadas parciales de $\hat{\nu}$ a lo largo de la diagonal $z = t\mathbf{1}$. Debido a que $\hat{\nu}$ es multilineal, estas derivadas son polinomios univariados de grado a lo sumo $n-1$. Por lo tanto, las atribuciones exactas pueden recuperarse mediante interpolación de Vandermonde a partir de un pequeño número de puntos de prueba, eliminando la varianza de Monte Carlo.

3. Sustituto de Red de Tensores Alineado al Grafo

Dado que el tensor de coalición completo $T$ tiene $2^n$ entradas, TN-SHAP-G aprende un sustitto de red de tensores alineado al grafo $\hat{\nu}$ para aproximar la extensión multilineal.

• Topología: La topología de la TN refleja el grafo de entrada $G$. Cada nodo $v$ corresponde a un tensor núcleo $A^{(v)}$ con un índice físico (que codifica la membresía de la coalición) e índices de enlace que conectan con sus vecinos.
• Expresividad: La dimensión de enlace $\chi$ controla la capacidad de flujo de información a través de los separadores del grafo. Los autores demuestran que para juegos con bajo rango de corte (cut-rank) a través de los separadores del grafo (una propiedad de muchos patrones de dependencia local en las GNN), esta estructura alineada al grafo logra una mayor expresividad con menos parámetros que las factorizaciones basadas en cadenas como los Tensor Trains (TT), que sufren de inflación de rango en grafos cíclicos o complejos.
• Entrenamiento: El sustituto se entrena mediante la minimización del riesgo empírico sobre un pequeño conjunto de coaliciones muestreadas ($M \ll 2^n$). Las garantías teóricas muestran que el presupuesto de consultas requerido escala de forma casi lineal con el número de nodos para grafos de grado acotado.

4. Inferencia

Una vez entrenado, el sustituto permite la computación determinista de:

• Valores de Shapley de primer orden: Mediante la integración 1D de la derivada de la diagonal.
• Interacciones de orden superior: Mediante derivadas parciales mixtas e integración.
  Todas las cantidades se computan a partir del único sustituto entrenado sin consultas adicionales al modelo de caja negra.

Contribuciones Clave

1. Representación Alineada al Grafo con Principios Fundamentados: Introducción de un sustituto de red de tensores que comprime la tabla de valores de coalición exponencial en una representación multilineal de bajo rango y alineada al grafo. Esto alinea el sesgo inductivo del sustituto con los patrones de dependencia local en los predictores de grafos.
2. Computación Determinista Amortizada: Aprovechamiento de la estructura multilineal para computar valores de Shapley e interacciones mediante interpolación de Vandermonde de forma cerrada. Esto requiere solo $O(n)$ evaluaciones del sustituto e introduce ninguna varianza de Monte Carlo. Un único sustituto amortiza la computación de todos los índices de primer y orden superior.
3. Garantías Teóricas:
   • Demostración de que el sustituto aprendido determina de forma única la extensión multilineal, asegurando la recuperación exacta de las atribuciones en el sustituto.
   • Caracterización de la expresividad mediante argumentos de rango de corte, demostrando que las TN alineadas al grafo superan a las factorizaciones de cadena (como TT) para juegos con estructura de grafo.
   • Límites de complejidad de consulta del profesor (teacher-query) casi lineales para grafos de grado acotado.
   • Garantías de estabilidad que vinculan el error de aproximación del juego con el error de Shapley.
4. Desempeño Empírico: Demostró una alta precisión en bancos de pruebas moleculares (Mutagenicity, Benzene), logrando una similitud de coseno $>0.99$ con los valores exactos de Shapley con tan solo 50 evaluaciones del modelo (10–100$\times$ menos que los baselines).

Resultados Experimentales

Los autores evaluaron TN-SHAP-G en bancos de pruebas moleculares (Mutagenicity, Benzene) y grafos de proteínas (PROTEINS).

• Corrección: En grafos pequeños ($n \le 20$), TN-SHAP-G coincide estrechamente con los valores exactos de Shapley (O1) e índices de interacción (O2), superando significativamente a GraphSVX y al muestreo por permutación tanto en precisión como en varianza.
• Eficiencia de Consulta: TN-SHAP-G logra una alta precisión de atribución con 10–100$\times$ menos consultas al modelo que los estimadores basados en muestreo (SHAP-IQ, GraphSVX).
• Escalabilidad: El método escala eficientemente a grafos con hasta 600 nodos utilizando una arquitectura de sustituto basada en el refinamiento (coarsening), mientras que la enumeración exacta y los métodos basados en muestreo se vuelven impracticables.
• Ablación de Estructura: Las TN alineadas al grafo superan consistentemente a los sustitutos de Tensor Train (TT) bajo condiciones de parámetros igualados, confirmando la ventaja teórica de las factorizaciones alineadas a la topología.
• Proxy de Fidelidad: Se observó una fuerte correlación entre el $R^2$ de prueba del sustituto y la similitud de coseno de los valores de Shapley recuperados, validando el $R^2$ como un proxy práctico para la precisión de la atribución.

Significación y Reivindicaciones

El artículo afirma que TN-SHAP-G hace avanzar el estado del arte en la explicabilidad de grafos al proporcionar un marco con principios fundamentales y agnóstico al modelo que supera la complejidad exponencial de la computación de Shapley. Al explotar la estructura de bajo rango inherente en los juegos cooperativos con estructura de grafo, el método permite la atribución determinista y libre de varianza con una complejidad de consulta casi lineal.

Los autores enfatizan que este enfoque democratiza el acceso a métodos de explicación rigurosos para profesionales con recursos computacionales limitados, particularmente en dominios de alto riesgo como el descubrimiento de fármacos y la predicción de propiedades moleculares. El marco es flexible, soportando varios semánticas de juego (por ejemplo, enmascaramiento de nodos, eliminación de nodos, valores de Myerson) y tipos de características, y se extiende más allá de las GNN de paso de mensajes estándar hacia otros predictores de grafos. El trabajo posiciona a las redes de tensores no solo como herramientas de compresión, sino como un mecanismo fundamental para el análisis tratable de juegos en datos estructurados.