Resumen Técnico: T3R - Adaptación de Prueba más Profunda para Redes Neuronales de Grafos mediante la Rotación de Gradientes

1. Planteamiento del Problema

Las Redes Neuronales de Grafos (GNN) desplegadas en sistemas del mundo real, como redes de distribución de agua, infraestructura 5G y redes informáticas, operan típicamente con pesos fijos tras el entrenamiento. Este paradigma estático conduce a una degradación del rendimiento al encontrar cambios en la distribución (deriva de concepto o concept drift), un suceso común en los datos de sensores donde las condiciones cambian con el tiempo.

Aunque el Entrenamiento en Tiempo de Prueba (TTT, por sus siglas en inglés) ofrece una solución al adaptar los modelos utilizando datos de prueba no etiquetados mediante una tarea auxiliar de Aprendizaje Auto-Supervisado (SSL), los enfoques existentes enfrentan limitaciones significativas en el dominio de los grafos:

1. Adaptación Superficial: Los métodos TTT convencionales (por ejemplo, Sun et al., 2020) típicamente actualizan solo el codificador compartido y el decodificador SSL durante la inferencia, dejando fijo el decodificador de la tarea principal. Esto restringe la capacidad del modelo para adaptarse plenamente a nuevas distribiones.
2. Desalineación de Tareas: Seleccionar una tarea SSL auxiliar que se alinee significativamente con una tarea objetivo específica basada en grafos es difícil. Los conflictos de gradiente entre tareas pueden obstaculizar la optimización conjunta.
3. Escasez de Datos: En dominios como los sistemas de agua, los datos etiquetados son escasos, y el reentrenamiento o el ajuste fino (fine-tuning) mediante simulaciones suele ser inviable debido a la falta de conocimiento completo del sistema.

2. Metodología: T3R

Los autores proponen T3R (Test-Time Training with layer-wise Rotograds), un nuevo marco que aprovecha la rotación de gradientes para permitir una adaptación más profunda en casi toda la arquitectura de la GNN sin necesidad de etiquetas de verdad fundamental (ground-truth).

Arquitectura Central

T3R utiliza una arquitectura en forma de Y que consta de:

• Un codificador compartido ($\theta_e$).
• Una cabeza principal ($\theta_m$) para la tarea objetivo.
• Una cabeza SSL auxiliar ($\theta_s$) para un objetivo sin etiquetas (por ejemplo, enmascaramiento de nodos).

Mecanismos Clave

1. Alineación de Gradientes y Afinidad de Tarea:
El método mide la afinidad de la tarea como la similitud de coseno entre los vectores de gradiente de las cabezas principal y SSL. Una alta alineación positiva indica que minimizar un objetivo también minimiza el otro. Los autores seleccionan empíricamente tareas SSL (específicamente el enmascaramiento de nodos) que exhiben la mayor alineación de gradiente positiva con la tarea principal.

2. Rotograds por Capa:
Para mitigar el conflicto de gradiente y mejorar la afinidad de la tarea, T3R integra Rotograd (Javaloy & Valera, 2022) en el marco de TTT.

• Fase de Entrenamiento: Se anteponen matrices de rotación ortogonales ($R^{(l)}_m, R^{(l)}_s$) a cada capa de los decodificadores principal y SSL. Se optimiza una pérdida de rotación ($L_{rot}$) de forma cíclica para maximizar la similitud de coseno entre los gradientes originales y los gradientes rotados y específicos de la tarea. Esto fomenta que los gradientes de ambas tareas se alineen en una dirección común.
• Fase de Tiempo de Prueba: Durante la inferencia, la pérdida de la tarea principal no está disponible. T3R utiliza las matrices de rotación aprendidas para generar gradientes sustitutos para la tarea principal. Específicamente, el gradiente SSL se rota de vuelta hacia la dirección media del gradiente original y se redirige para actualizar las capas del decodificador principal.
  • Fórmula: $\tilde{g}^{(l+1)}_m = R^{(l)\top}_s g^{(l)}_s R^{(l)}_m$.
  • Este mecanismo permite que toda la arquitectura (codificador compartido, decodificador SSL y casi todas las capas del decodificador principal) sea actualizada, en lugar de solo el codificador.

3. Desplazamientos de Hiperparámetros:
Los autores observan que el ajuste de los hiperparámetros durante la fase de tiempo de prueba mejora el rendimiento. Específicamente, aumentan el coeficiente de la pérdida auxiliar ($\lambda_{test}$) y reducen la intensidad de la aumentación (por ejemplo, tasas de enmascaramiento de nodos más bajas) para amplificar las señales SSL y evitar fallos en lotes (batches) de prueba pequeños.

3. Contribuciones Principales

El artículo detalla tres contribuciones primarias:

1. Método T3R: Un nuevo enfoque de adaptación que utiliza matrices de rotación por capa para mejorar la afinidad de la tarea durante el entrenamiento y generar gradientes sustitutos para una adaptación profunda y en toda la arquitectura durante el tiempo de prueba.
2. Selección Empírica de SSL: Una investigación sobre la selección de objetivos SSL basados en grafos que se alineen con las tareas principales, demostrando que la alineación de gradiente positiva se correlaciona con una mejor adaptabilidad.
3. Evaluación Exhaustiva: Experimentos integrales en conjuntos de datos de regresión (redes de distribución de agua) y clasificación (grafos moleculares) bajo configuraciones de Fuera de la Distribución (OOD) y de dominio cruzado.

4. Resultados Experimentales

Los autores evaluaron T3R en el conjunto de datos DiTEC-WDN (DWD) (regresión) y en los conjuntos de datos Open Graph Benchmark (OGB) (clasificación).

Regresión (Redes de Distribución de Agua)

• Rendimiento: T3R logró una mejora relativa del 38.38% en el RMSE en comparación con la Minimización del Riesgo Empírico (ERM) estándar y superó a otros modelos de base de adaptación (Tent, TTT estándar, TTT+Rotograd) en múltiples métricas (MAE, R2, PCC).
• Observación: El TTT estándar y TTT+Rotograd (que solo actualiza el codificador) mostraron ganancias limitadas o tuvieron un rendimiento inferior al ERM en este entorno, lo que resalta la necesidad de actualizar el decodificador principal.
• Validación Cruzada: En pruebas de dominio cruzado (entrenamiento en una topología de red, prueba en otra), T3R con un solo paso de adaptación ($t=1$) produjo los mejores resultados. Una adaptación más profunda (más pasos) generalmente degradó el rendimiento en las tareas de regresión.

Clasificación (Grafos Moleculares)

• Adaptación de Dominio Cruzado: En experimentos de adaptación entre ogbg-molbace y ogbg-molbbbp, T3R logró mejoras relativas significativas en ROC-AUC, alcanzando hasta un 69.56% de mejora relativa sobre la inferencia estándar sin adaptación.
• Comparación: Mientras que otros métodos mostraron mejoras leves o convergencia temprana, T3R demostró ganancias robustas, particularmente cuando los pasos de adaptación se ajustaron cuidadosamente.

Estudios de Ablación

• Capas de Rotación: Eliminar las capas de rotación aumentó la incertidumbre y degradó el rendimiento, confirmando la necesidad de la adaptación por capa.
• Coeficiente de Pérdida: El ajuste del coeficiente de la pérdida auxiliar en tiempo de prueba ($\lambda_{test}$) para que fuera mayor que el coeficiente de entrenamiento mejoró el rendimiento de la tarea principal, mientras que los valores excesivos provocaron una sobre-adaptación.

5. Significado y Reivindicaciones

El artículo afirma que T3R aborda la limitación crítica de los métodos TTT existentes al permitir una adaptación más profunda en toda la arquitectura del modelo, no solo en el codificador. Al aprovechar la rotación de gradientes para crear gradientes sustitutos, T3R permite que los modelos se adapten eficazmente en escenarios donde:

• Los datos etiquetados no están disponibles o son costosos de recolectar.
• Los cambios en la distribución son frecuentes (por ejemplo, la deriva de sensores en la infraestructura).
• El ajuste fino o el reentrenamiento convencional es inviable.

Los autores posicionan a T3R como un paso hacia el desarrollo de procesos de adaptación robustos y sin etiquetas para sistemas basados en grafos en tiempo real, particularmente en dominios de infraestructura crítica como las redes de agua. Reconocen que, si bien el método es prometedor, la relación entre la alineación de gradientes y la adaptabilidad requiere una justificación teórica adicional, y el costo computacional de los pasos de actualización de gradiente adicionales es un compromiso (trade-off) para la mejora de la precisión.