How to Make LMs Strong Node Classifiers?

Resumen Técnico: AUGLM – Modelo de Lenguaje de Grafos Aumentado para la Clasificación de Nodos

Planteamiento del Problema
Si bien los Modelos de Lenguaje (LM) han mostrado potencial en tareas de aprendizaje en grafos, los enfoques existentes para la clasificación de nodos en Grafos con Atributos de Texto (TAGs) enfrentan un compromiso fundamental entre el rendimiento y la flexibilidad arquitectónica. Los métodos de vanguardia (SOTA) como InstructGLM y LLaGA logran resultados sólidos modificando la arquitectura del LM —específicamente, expandiendo el vocabulario con tokens únicos para cada nodo o inyectando incrustaciones (embeddings) conscientes de la topología en el espacio de tokens—. Aunque efectivos, estas modificaciones comprometen la generalidad del LM base, impidiendo un aprendizaje multitarea fluido entre diversos conjuntos de datos y restringiendo el uso de modelos propietarios o de código cerrado donde el acceso a la arquitectura no está disponible. La pregunta central que aborda este trabajo es: ¿Cómo pueden los LMs de instrucción texto-a-texto, listos para usar (off-the-shelf), lograr un rendimiento competitivo en la clasificación de nodos sin ninguna modificación arquitectónica?

Metodología: AUGLM
Los autores proponen AUGLM (Modelo de Lenguaje de Grafos Aumentado), un marco que mejora los LMs estándar mediante estrategias de aumentación de datos en lugar de cambios arquitectónicos. AUGLM reformula la clasificación de nodos como una tarea de texto a texto, preservando la estructura original del LM mientras enriquece el contexto de entrada. El marco se basa en tres componentes principales:

1. Agregación Basada en Recuperación:
   En lugar de depender de una simple agregación de vecinos de múltiples saltos (que puede introducir ruido o causar una "explosión de vecinos"), AUGLM emplea dos estrategias de recuperación complementarias para proporcionar información contextual rica:

   • Recuperación Topológica: Utiliza el PageRank Personalizado (PPR) para identificar a los vecinos más relevantes para un nodo objetivo basándose en la estructura del grafo, en lugar de distancias de saltos fijos. Las características textuales de estos K-vecinos superiores de PPR se concatenan para formar un contexto estructural.
   • Recuperación Semántica Prototípica: Inspirado en la Generación Aumentada por Recuperación (RAG) y los Transformadores de Grafos, este módulo trata las características textuales de todos los nodos que no son el objetivo como un corpus externo. Recupera "prototipos" —ejemplos representativos para cada clase— mediante el entrenamiento previo de una GNN ligera para generar predicciones de nodos. Se seleccionan como prototipos los ejemplos con mayor confianza para cada clase. El sistema luego recupera el prototipo semánticamente más similar (y sus vecinos de PPR asociados) al texto del nodo objetivo.

2. Guía del Clasificador (Poda de Candidatos):
   Para abordar la limitada comprensión inherente de la topología de grafos del LM, AUGLM incorpora la guía de un modelo especialista. Una GNN ligera preentrenada predice un conjunto de $I$ etiquetas candidatas para cada nodo. Estos candidatos se convierten en texto y se incluyen en el prompt. Esta estrategia efectivamente poda el espacio de búsqueda, permitiendo que el LM se concentre su razonamiento en distinguir entre candidatos altamente relevantes e informados por la topología, en lugar de todo el espacio de etiquetas.

3. Entrenamiento y Optimización:
   El marco involucra tres módulos parametrizados: una GNN ligera ($\psi$) para generar prototipos y candidatos, un recuperador semántico ($\phi$) y el LM base ($\theta$).

   • La GNN se preentrena de forma independiente.
   • El recuperador semántico se entrena utilizando una pérdida de correspondencia de distribución (divergencia KL) para alinear las probabilidades de recuperación con la probabilidad del LM de generar la etiqueta de verdad de campo (ground-truth) correcta.
   • El LM base se ajusta (fine-tuning) utilizando una pérdida estándar de log-verosimilitud negativa (NLL) por token sobre los prompts de entrada aumentados.
   • Crucialmente, los grafos computacionales del recuperador y del LM se desacoplan mediante operadores de parada de gradiente (stop-gradient) durante fases específicas del entrenamiento, asegurando una retropropagación eficiente.

Contribuciones Clave

• Preservación Arquitectónica: A diferencia de los métodos SOTA anteriores, AUGLM no requiere modificaciones en la arquitectura o el vocabulario del LM, lo que permite el uso de cualquier LM de texto a texto listo para usar.
• Aumentación de Estrategia Dual: La introducción de un mecanismo de recuperación híbrido que combina información topológica (PPR) y semántica (basada en prototipos) para superar las limitaciones de la simple agregación de vecinos.
• Guía Ligera: Un método novedoso para transferir el sesgo inductivo topológico de una GNN a un LM mediante la poda de etiquetas candidatas, mejorando el rendimiento sin alterar los parámetros internos del LM.
• Flexibilidad Multitarea: El marco soporta el entrenamiento conjunto en múltiples y diversos conjuntos de datos sin la necesidad de expansiones de tokens específicas de cada conjunto, una capacidad que suele perderse en los enfoques que modifican la arquitectura.

Resultados Experimentales
Los autores evaluaron AUGLM en cuatro conjuntos de datos de TAG del mundo real: Cora, Pubmed, ogbn-arxiv y ogbn-products.

• Rendimiento: Equipado con backbones Flan-T5, AUGLM superó consistentemente a los clasificadores de nodos de salida de texto SOTA (incluyendo InstructGLM y LLaGA).
• Comparación con Modelos de Salida de Vectores: AUGLM logró un rendimiento comparable a, y en algunos casos superando (por ejemplo, en ogbn-products), a los mejores modelos de salida de vectores (como TAPE y GLEM+RevGAT), a pesar de operar puramente en el dominio del texto.
• Estudios de Ablación: Los experimentos confirmaron que todos los tres componentes (recuperación topológica, recuperación semántica y poda de etiquetas) contribuyen positivamente al rendimiento, aunque su importancia relativa varía según el conjunto de datos.
• Entrenamiento Conjunto: Un modelo entrenado conjuntamente en los cuatro conjuntos de datos alcanzó un rendimiento comparable al de los modelos entrenados por separado en cada uno, demostrando la capacidad del marco para generalizar a través de diversas estructuras de grafos sin degradación del rendimiento.

Significancia y Reivindicaciones
El artículo posiciona a AUGLM como un paso crucial hacia el cierre de la brecha entre los clasificadores de nodos especializados y de tareas específicas y los modelos fundacionales de propósito general. Al demostrar que se puede lograr un alto rendimiento en el aprendizaje de grafos mediante la aumentación de la entrada en lugar de la modificación arquitectónica, el trabajo sugiere un camino hacia modelos de aprendizaje de grafos más versátiles y ampliamente aplicables. Los autores afirman que este enfoque retiene los beneficios del ajuste de instrucción (instruction tuning) —específicamente la capacidad de entrenar conjuntamente en diversos conjuntos de datos— mientras supera la inflexibilidad de los métodos previos de integración de grafos-LM. El trabajo concluye que, si bien los LMs actuales tienen limitaciones inherentes para comprender las estructuras de los grafos, estas pueden mitigarse eficazmente mediante la aumentación estratégica de datos, allanando el camino para modelos unificados que sobresalen en múltiples tareas de aprendizaje de grafos.