Resumen Técnico: El Grafo es una Regularización Natural: Revisitando la Cuantización de Vectores para el Aprendizaje de Representaciones de Grafos

1. Declaración del Problema

La Cuantización de Vectores (VQ) ha emergido como una técnica poderosa para el aprendizaje de representaciones discretas y comprimidas en los dominios de visión y lenguaje, ofreciendo un camino hacia Modelos Fundacionales de Grafos (GFM) escalables mediante la tokenización de estructuras de grafos. Sin embargo, este artículo identifica un desafío fundamental y poco explorado específico del dominio de los grafos: el colapso del libro de códigos (codebook collapse).

El colapso del libro de códigos ocurre cuando la mayoría de los puntos de datos de entrada se mapean a solo unos pocos vectores de código (codewords) en el libro de códigos aprendible, dejando la gran mayoría del libro de códigos infrautilizada. Los autores observan que este fenómeno es sistemático en el VQ de grafos, ocurriendo incluso al aplicar estrategias de mitigación de vanguardia (SOTA) desarrolladas para visión y lenguaje (por ejemplo, Media Móvil Exponencial, reinicio del libro de códigos, codificadores preentrenados). Este colapso limita severamente la expresividad y la generalización de los tokens de grafos, lo que conduce a representaciones toscas y un rendimiento degradado en tareas posteriores.

El artículo diagnostica las causas raíz de este colapso desde dos perspectivas:

1. Perspectiva de los Datos: Los datos de grafos poseen propiedades inherentes como la redundancia de características y la alta densidad de conectividad. El estudio encuentra una correlación positiva entre la redundancia de características (medida por PCA@95%) y la severidad del colapso, y una correlación negativa entre la densidad de conectividad (grado promedio de los nodos) y la utilización del libro de códigos. La conectividad densa y las características redundantes exacerban la tendencia de los nodos a mapearse a tokens idénticos.
2. Perspectiva de la Optimización: El VQ determinista depende de la asignación dura (hard assignment) basada en la proximidad (selección del vecino más cercano). Esto crea un bucle de retroalimentación autorreforzado: los codewords seleccionados con frecuencia reciben más actualizaciones de gradiente y se vuelven cada vez más dominantes, mientras que los codewords raramente seleccionados no reciben actualizaciones y permanecen inactivos. Este "efecto capullo" (cocoon effect) bloquea al codificador en un conjunto estrecho de tokens, impidiendo la exploración de la capacidad total del libro de códigos.

2. Metodología: RGVQ

Para abordar estos problemas, los autores proponen RGVQ (Regularized Graph Vector Quantization), un marco que integra la topología del grafo y la similitud de características como señales de regularización explícitas. El método consta de dos componentes principales:

A. Reparametrización Gumbel-Softmax

Para romper el bucle de autorrefuerzo de la asignación dura determinista, RGVQ reemplaza la operación argmax no diferenciable con una reparametrización Gumbel-Softmax.

• En lugar de seleccionar un único codeword más cercano, el modelo calcula los logits de asignación basados en la distancia entre los embeddings de los nodos y todos los codewords.
• Estos logits se perturban con ruido Gumbel y se pasan a través de una función softmax controlada por temperatura para producir una distribución de probabilidad diferenciable sobre el libro de códigos.
• El embedding cuantizado se calcula como un promedio ponderado de todos los codewords basado en estas probabilidades.
• Efecto: Esto asegura que los gradientes fluyan hacia todos los codewords proporcionalmente a sus probabilidades de asignación, no solo hacia el seleccionado, permitiendo que los codewords inactivos reciban actualizaciones y sean reactivados.

B. Regularización Contrastiva Sensible a la Estructura

Para mitigar los efectos de la redundancia de los datos de grafos y la conectividad densa, RGVQ introduce un término de regularización que penaliza explícitamente el uso de tokens excesivamente concentrados entre nodos disímiles.

• Conjunto Positivo ($N_P$): Para un nodo ancla dado, el conjunto positivo incluye nodos que están estructuralmente conectados (vecinos) o son semánticamente similares (top-$K$ nodos con características similares).
• Conjunto Negativo ($N_N$): Incluye nodos que no están conectados ni son similares en características.
• Función de Pérdida: El modelo emplea una pérdida InfoNCE para fomentar que los nodos en el conjunto positivo compartan distribuciones de asignación de tokens similares, mientras que penaliza a los nodos en el conjunto negativo por tener distribuciones superpuestas.
• Efecto: Esta regularización desincentiva que el modelo mapee nodos estructural o semánticamente distintos al mismo token, promoviendo así la diversidad de los tokens y evitando el colapso inducido por la redundancia de características.

La función de pérdida total combina la pérdida de reconstrucción estándar (reconstrucción de características y de enlaces), la pérdida del libro de códigos, la pérdida de compromiso (commitment loss) y la nueva pérdida de regularización sensible a la estructura.

3. Contribuciones Clave

El artículo realiza las siguientes contribuciones:

1. Estudio Empírico Sistemático: Proporciona la primera evaluación sistemática del colapso del libro de códigos en el VQ de grafos, demostrando que las estrategias de mitigación estándar de otros dominios fallan al resolver el problema en los datos de grafos.
2. Diagnóstico del Colapso: Identifica que el colapso del libro de códigos en grafos es impulsado por las propiedades intrínsecas de los datos (redundancia y conectividad) y la dinámica de optimización de la asignación dura determinista.
3. Marco RGVQ: Propone un nuevo marco diferenciable que combina la reparametrización Gumbel-Softmax con la regularización contrastiva sensible a la estructura para mejorar la utilización del libro de códigos.
4. Evaluación Exhaustiva: Valida RGVQ a través de múltiples arquitecturas de base de VQ de grafos y tareas posteriores, mostrando mejoras consistentes tanto en la utilización del libro de códigos como en el rendimiento de las tareas.

4. Resultados Experimentales

Los autores realizaron experimentos extensos en varios conjuntos de datos de grafos homofílicos y heterofílicos (por ejemplo, Cora, PubMed, WikiCS, Amazon-Ratings, Roman-Empire).

• Utilización del Libro de Códigos: RGVQ supera significativamente al VQ de grafos vanilla y a sus variantes (incluyendo EMA, parámetros afines, reinicio del libro de códigos y codificadores preentrenados) en términos de perplejidad del libro de códigos. Por ejemplo, en el conjunto de datos PubMed, mientras que el VQ vanilla logró una perplejidad de ~4.14 (indicando un colapso severo), RGVQ logró ~319.09, acercándose al máximo teórico. RGVQ mantuvo una alta utilización incluso a medida que aumentaban los tamaños del libro de códigos.
• Transferibilidad: Al integrarse en el Modelo Fundacional de Grafos (GFT), RGVQ alcanzó un rendimiento de vanguardia (SOTA) a través de evaluaciones de tareas cruzadas y dominios cruzados (clasificación de nodos, predicción de enlaces, clasificación de grafos), superando tanto a las GNN supervisadas como a otros métodos de aprendizaje autosupervisado.
• Serialización: En el contexto de los Transformadores de Grafos Cuantizados (GQT), los modelos mejorados con RGVQ mejoraron consistentemente la precisión de la clasificación de nodos en comparación con GQT con otros métodos anti-colapso.
• Estudios de Ablación:
  • Eliminar la reparametrización Gumbel-Softmax o la regularización sensible a la estructura provocó una caída significativa en el rendimiento, confirmando que ambos componentes son mutuamente dependientes para prevenir el colapso.
  • El método fue robusto a variaciones en la temperatura ($\tau$) y en el número de muestras contrastivas, demostrando que temperaturas bajas y fijas son suficientes.
  • RGVQ demostró robustez a través de varios niveles de heterofilia de grafos.

5. Significado y Reivindicaciones

El artículo afirma que el colapso del libro de códigos no es un problema incidental, sino un cuello de botella fundamental en el aprendizaje de tokens de grafos discretos. Al tratar la estructura del grafo como un regularizador natural, RGVQ ofrece una solución de tipo "plug-in" flexible que mejora la expresividad y la transferibilidad de los tokens de grafos.

Los autores enfatizan que su trabajo no propone una nueva arquitectura de modelo generativo per se, sino más bien una mejora metodológica al proceso de cuantización que permite un aprendizaje más efectivo de las representaciones de grafos. El significado radica en permitir la construcción de Modelos Fundacionales de Grafos más robustos al asegurar que el vocabulario de tokens aprendido sea plenamente utilizado, capturando así patrones estructurales y semánticos más finos esenciales para las tareas posteriores. El artículo concluye que RGVQ proporciona un marco flexible que mantiene la diversidad de los tokens y escala eficazmente a tamaños de libro de códigos más grandes, abordando una limitación crítica en el aprendizaje actual de representaciones de grafos.