Resumen Técnico: PreLort – LoRA de Prefijo Anidado para el Ajuste Fino Federado bajo Heterogeneidad de Rango

1. Planteamiento del Problema

El Aprendizaje Federado (FL) permite el entrenamiento colaborativo de Modelos de Lenguaje de Gran Escala (LLMs) sin compartir datos brutos, pero enfrenta desafíos significativos al adaptar modelos fundacionales mediante métodos de ajuste fino eficiente en parámetros (PEFT) como Low-Rank Adaptation (LoRA). Un obstáculo principal es la heterogeneidad de rango: en entornos realistas de dispositivos cruzados, los clientes poseen diversas restricciones computacionales y de memoria, lo que los lleva a inicializar adaptadores LoRA con diferentes rangos ($r$).

Las soluciones existentes para el LoRA federado heterogéneo se centran principalmente en la fase de agregación, empleando a menudo el relleno con ceros (zero-padding) para alinear adaptadores de diferentes rangos antes del promedio. Sin embargo, los autores identifican una limitación fundamental en estos enfoques: la desalineación de representaciones. Incluso cuando la agregación está matemáticamente definida, combinar modelos entrenados de forma independiente con diferentes rangos conduce a inconsistencias. Específicamente, los clientes de menor rango no se benefician de las actualizaciones de los clientes de mayor rango porque la información no se distribuye jerárquicamente a través de las dimensiones de rango. El relleno con ceros además diluye las contribuciones de los clientes de alto rango al introducir segmentos con valores de cero, lo que resulta en un uso subóptimo de las representaciones de bajo rango compartidas.

2. Metodología: PreLort

Los autores proponen PreLort (Prefix-Nested LoRA), un marco que aborda la heterogeneidad de rango estructurando conjuntamente la dinámica de entrenamiento y las reglas de agregación. La filosofía central es organizar las dimensiones del adaptador en una jerarquía de prefijos, asegurando que las dimensiones de menor rango codifiquen información relevante para la tarea compartida entre todos los clientes, mientras que las dimensiones de mayor rango capturan la capacidad incremental.

2.1 Estructura de Bajo Rango Anidada

PreLort impone una organización de estructura de prefijo en las dimensiones del adaptador. Para un cliente con rango $r_k$, solo las primeras $r_k$ dimensiones están activas. Crucialmente, durante el entrenamiento local, el adaptador no se optimiza únicamente en su rango completo. En su lugar, se optimiza bajo múltiples configuraciones truncadas (sub-rangos de prefijo).

• Mecanismo: Un cliente con rango $r_k$ es entrenado para desempeñarse bien no solo en $r_k$, sino también en todos los niveles de rango inferiores presentes en la federación ($r \in \{r \mid r \le r_k\}$).
• Objetivo: Esto obliga al modelo a concentrar las señales relevantes de la tarea en las dimensiones principales (de bajo rango), asegurando que estas dimensiones sigan siendo informativas y útiles incluso cuando se truncan.

2.2 Estrategia de Entrenamiento Local

El objetivo local para un cliente $k$ se define como el promedio de las pérdidas de entropía cruzada computadas sobre todos los sub-rangos de prefijo válidos $R_k$:
$$\mathcal{L}_k = \frac{1}{|R_k|} \sum_{r \in R_k} \mathcal{L}_{CE}(f(\cdot; \theta, r), \mathcal{D}_k)$$
En la práctica, cada paso de entrenamiento implica una pasada de forward/backward de rango completo y pasadas adicionales para todos los sub-rangos. Esta estrategia asegura que los segmentos de prefijo se optimicen explícitamente a la misma granularidad utilizada durante la agregación, fomentando una jerarquía consistente donde las dimensiones inferiores capturan la información más crítica.

2.3 Agregación por Segmentos

Para evitar la dilución causada por el relleno con ceros, PreLort introduce una regla de agregación por segmentos que opera en el espacio del gradiente (actualización).

• Proceso: En lugar de promediar todos los clientes para cada dimensión, el servidor agrega los gradientes para cada segmento de rango $[r_{l-1}, r_l)$ solo sobre el subconjunto de clientes cuyo rango $r_k$ es suficiente para contribuir a ese segmento (es decir, $r_k \ge r_l$).
• Fórmula: Para el $l$-ésimo segmento, el gradiente agregado es:
  $$\nabla A_{global}^{r_{l-1}:r_l} = \frac{\sum_{k: r_k \ge r_l} n_k \nabla A_k^{r_{l-1}:r_l}}{\sum_{k: r_k \ge r_l} n_k}$$
• Beneficio: Esto evita la dilución de la señal de los clientes de menor rango con relleno de ceros y asegura que las dimensiones de mayor rango sean actualizadas solo por los clientes que realmente poseen la capacidad de aprenderlas, mientras que las dimensiones inferiores se benefician de la población completa.

3. Contribuciones Clave

El artículo describe tres contribuciones primarias:

1. Identificación de la Desalineación de Representaciones: Los autores señalan que la falta de alineación estructurada a través de las dimensiones de rango es una limitación clave en los métodos actuales de LoRA federado heterogéneo, lo que conduce a representaciones inconsistentes.
2. Formulación Anidada y Estrategia Unificada: Proponen una formulación de bajo rango anidada que impone una estructura de prefijo común. Esto se combina con una estrategia de entrenamiento y agregación unificada que permite una agregación consistente a través de rangos heterogéneos sin necesidad de relleno con ceros.
3. Validación Empírica: Se demuestra que el método mejora consistentemente la calidad y estabilidad del modelo sobre los métodos previos de LoRA federado heterogéneo en tareas de ajuste de instrucciones (instruction-tuning).

4. Resultados Experimentales

Los autores evaluaron PreLort en dos modelos base (TinyLlama-1.1B y Qwen2.5-0.5B) a través de tres conjuntos de datos: Alpaca, Databricks-dolly-15 y 20 Newsgroups.

• Métricas de Rendimiento: PreLort superó consistentemente a las líneas base (ZeroPad, HetLoRA, FLoRA) en Exactitud (MMLU, 20 Newsgroups) y puntuaciones ROUGE-L.
  • En TinyLlama-1.1B, PreLort alcanzó un 33.86% de exactitud en MMLU (Dolly) y un 30.76% en Alpaca, superando a las mejores líneas base.
  • En Qwen2.5-0.5B, alcanzó un 39.85% de exactitud en MMLU (Dolly).
• Perplejidad: PreLort logró la menor perplejidad en Dolly y NewsGroup para ambos modelos base, demostrando una verosimilitud de tokens superior.
• Robustez de Rango: El análisis mostró que PreLort mantiene un rendimiento estable a través de diferentes niveles de truncamiento de prefijo. En contraste, las líneas base como "Nested Aggregation Only" exhibieron inestabilidad, particularmente en rangos altos, confirmando que la agregación por sí sola es insuficiente sin el objetivo de entrenamiento anidado.
• Importancia por Dimensión: La visualización de la importancia de las dimensiones reveló que PreLort concentra la señal relevante de la tarea en las dimensiones principales (una distribución estructurada), mientras que el entrenamiento estándar produce una distribución uniforme. Esta priorización jerárquica facilita actualizaciones estables a través de rangos heterogéneos.

5. Significado y Reivindicaciones

El artículo sostiene que PreLort representa un cambio en el abordaje de la heterogeneidad de rango: en lugar de tratarlo como un problema de agregación post-hoc, se resuelve estructurando el proceso de entrenamiento mismo.

• Argumento Central: Los autores argumentan que abordar la heterogeneidad de rango requiere no solo una mejor agregación, sino también una alineación estructurada de representaciones a través de los rangos.
• Eficiencia: El método permite que los clientes de bajo rango se beneficien de la información más rica aportada por los clientes de mayor rango porque las dimensiones de prefijo se aprenden y agregan de manera consistente.
• Modestia sobre las Limitaciones: Los autores reconocen que la estrategia de entrenamiento anidado introduce aproximadamente un 40% de sobrecarga computacional adicional debido a las múltiples pasadas de forward/backward por lote. También señalan que la efectividad depende de suficientes pasos de optimización local; con épocas locales muy limitadas, el modelo puede no aprender completamente representaciones de prefijo de bajo rango bien alineadas.

En conclusión, PreLort demuestra que la heterogeneidad de rango en el LoRA federado se maneja mejor estructurando conjuntamente cómo se entrenan los adaptadores (mediante la optimización de prefijo anidado) y cómo se agregan sus actualizaciones (mediante reglas de segmentación), resultando en representaciones de prefijo alineadas y robustas que superan a los métodos de agregación heterogénea existentes.