Resumen Técnico: LoRIF – Funciones de Influencia de Bajo Rango para la Atribución de Datos de Entrenamiento Escalable

Declaración del Problema

La Atribución de Datos de Entrenamiento (TDA) tiene como objetivo identificar qué ejemplos de entrenamiento específicos influyeron más en la predicción de un modelo sobre una consulta dada. Aunque las Funciones de Influencia (IF) proporcionan un enfoque con fundamento teórico para la TDA al aproximar el reentrenamiento leave-one-out mediante gradientes y matrices hessianas inversas, han permanecido efectivamente inaccesibles para modelos de escala moderna (miles de millones de parámetros) entrenados en millones de ejemplos.

Las variantes escalables existentes, como TRAK, LoGRA y TrackStar, abordan los cuellos de botella computacionales almacenando gradientes proyectados por ejemplo y utilizando una aproximación de Hessiano de Gauss-Newton amortiguada. Sin embargo, estos métodos enfrentan dos cuellos de botella críticos de escalabilidad que crean una compensación entre la calidad de la atribución y la viabilidad computacional:

1. Almacenamiento de Gradientes y E/S: Almacenar los gradientes proyectados por ejemplo para $N$ ejemplos de entrenamiento requiere un espacio de $O(ND)$, donde $D$ es la dimensión de proyección efectiva. En el momento de la consulta, cargar estos gradientes domina la latencia (aproximadamente el 96% del tiempo de consulta).
2. Escala de la Hessiana Inversa: Formar y almacenar la aproximación de la Hessiana inversa de $D \times D$ requiere una memoria de $O(D^2)$.

Aumentar la dimensión de proyección $D$ es necesario para mejorar la calidad de la atribución, pero hacerlo exacerba tanto los cuellos de botella de almacenamiento/E/S como los de memoria, haciendo que la atribución de alta calidad sea poco práctica para modelos de escala fronteriza.

Metodología: LoRIF

LoRIF (Funciones de Influencia de Bajo Rango) aborda estos cuellos de botella explotando la estructura de bajo rango inherente en los gradientes de las redes neuronales. El método sigue una tubería de tres pasos:

1. Almacenamiento de Gradientes de Bajo Rango (Abordando el Cuello de Botella de E/S)

En lugar de almacenar la matriz completa de gradientes proyectados por ejemplo $\tilde{G}_i \in \mathbb{R}^{d_1 \times d_2}$ (donde $D = d_1 d_2$), LoRIF almacena una factorización de rango-$c$:
$$\tilde{G}_i \approx u_i v_i^\top$$
donde $u_i \in \mathbb{R}^{d_1 \times c}$ y $v_i \in \mathbb{R}^{d_2 \times c}$.

• Mecanismo: Esta factorización se calcula mediante iteraciones de potencia por bloques.
• Impacto: El almacenamiento y la E/S en tiempo de consulta por capa por muestra se reducen de $O(D)$ a $O(c(\sqrt{D}))$. Los resultados empíricos sugieren que incluso con $c=1$, la aproximación retiene una calidad de atribución significativa, y priorizar un $D$ mayor sobre un $c$ más alto produce una mejor relación costo-eficiencia.

2. Aproximación de Curvatura mediante SVD Truncado (Abordando el Cuello de Botella de Memoria)

Para evitar formar la Hessiana inversa densa de $D \times D$ $(G^\top G + \lambda I)^{-1}$, LoRIF calcula una Descomposición en Valores Singulares (SVD) truncada de rango-$r$ de la matriz de gradiente agregada $G$:
$$G \approx U_r \Sigma_r V_r^\top$$
Utilizando la identidad de la matriz de Woodbury, el término de la Hessiana inversa se aproxima en un subespacio de dimensión $r$:
$$(V_r \Sigma_r^2 V_r^\top + \lambda I)^{-1} = \frac{1}{\lambda}I - \frac{1}{\lambda^2} V_r \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} V_r^\top$$

• Mecanismo: Solo se almacenan los $r$ vectores singulares superiores ($V_r$) y los valores singulares ($\Sigma_r$).
• Impacto: La complejidad de memoria para el término de curvatura se reduce de $O(D^2)$ a $O(Dr)$. Esto permite una dimensión de proyección efectiva $D$ significativamente mayor dentro de un presupuesto de memoria fijo.

3. Cálculo de la Puntuación de Influencia

En el momento de la consulta, las puntuaciones de influencia se calculan proyectando el gradiente de la consulta en el subespacio de dimensión $r$ y utilizando los factores de bajo rango almacenados para los gradientes de entrenamiento. La puntuación de influencia $I(x_{tr}, x_{te})$ se aproxima como:
$$I(x_{tr}, x_{te}) \approx \frac{1}{\lambda} g_{te}^\top g_{tr} - \frac{1}{\lambda^2} (g'_{te})^\top \left( \Sigma_r^{-2} + \frac{1}{\lambda}I_r \right)^{-1} g'_{tr}$$
donde $g'$ representa los gradientes proyectados sobre los vectores singulares $V_r$. El primer término se calcula eficientemente utilizando los factores de rango-$c$, y el segundo término implica solo operaciones de $O(r)$ debido a la naturaleza diagonal de la matriz intermedia.

Contribuciones Clave

1. Identificación de Cuellos de Botella: El artículo identifica que la dimensión de proyección $D$ es la restricción principal en la TDA estilo IF escalable, creando una compensación calidad-escalabilidad donde la atribución precisa (gran $D$) es actualmente poco práctica.
2. Factorización de Rango-$c$: La propuesta de almacenar los gradientes proyectados como factores de rango-$c$, reduciendo el almacenamiento por ejemplo por capa de $O(D)$ a $O(c\sqrt{D})$.
3. Aproximación de Curvatura de Rango-$r$: La propuesta de utilizar SVD truncado y la identidad de Woodbury para aproximar el término de curvatura, reduciendo la memoria por capa de $O(D^2)$ a $O(Dr)$.
4. Demostración de Escalabilidad: La primera demostración de atribución de datos efectiva en modelos de hasta 70 mil millones de parámetros (Apertus-70B) con millones de ejemplos de entrenamiento, manteniendo costos y latencia prácticos.

Resultados Experimentales

LoRIF fue evaluado en modelos que van desde 124M hasta 70B de parámetros (GPT2-small, OLMo-3-7B, Apertus-70B) con conjuntos de datos que contienen hasta 3.8 millones de ejemplos.

• Almacenamiento y Latencia: En comparación con LoGRA, LoRIF logra un almacenamiento de gradientes 2.3–20× menor y una latencia de consulta 1.3–20× menor mientras iguala o supera la calidad de atribución.
  • En GPT2-small, LoRIF redujo el almacenamiento hasta en 9 veces y la latencia en 5.6 veces en regímenes de almacenamiento medio.
  • En OLMo-3-7B, LoRIF logró puntuaciones de parche de cola casi idénticas a las de LoGRA con 20.3× menos almacenamiento y 22.0× menos latencia.
  • En Apertus-70B, LoRIF mejoró las puntuaciones de parche de cola en 1.5–2.2× sobre LoGRA con menor almacenamiento y latencia.
• Calidad de Atribución: Medida mediante la Puntuación de Modelado de Datos Lineal (LDS) y la Puntuación de Parche de Cola, LoRIF mejoró consistentemente la frontera de Pareto. Al permitir dimensiones de proyección efectivas $D$ más grandes (que anteriormente eran inviables debido a restricciones de memoria), LoRIF logró una mayor calidad que las líneas base restringidas a un $D$ más pequeño.
• Evaluación Cualitativa: Una evaluación con LLM como juez en los ejemplos recuperados del top-1 mostró que LoRIF recupera ejemplos con mayor relevancia semántica (puntuación promedio 3.04 vs. 2.49 en OLMo-3-7B) y menos recuperaciones completamente irrelevantes en comparación con LoGRA.

Significado y Afirmaciones

El artículo afirma que LoRIF hace que la Atribución de Datos de Entrenamiento basada en gradientes sea práctica a escala fronteriza. Al explotar la estructura de bajo rango de los gradientes, LoRIF resuelve la compensación calidad-escalabilidad que ha obstaculizado la aplicación de funciones de influencia a los grandes modelos de lenguaje.

Los autores enfatizan que LoRIF no depende de nuevos primitivos de bajo rango, sino que demuestra cómo las estructuras de bajo rango existentes (observadas tanto en gradientes por ejemplo como agregados) pueden explotarse sistemáticamente en métodos de influencia de gradientes proyectados. Los resultados sugieren que la TDA basada en gradientes ahora puede utilizarse para la auditoría, depuración y verificaciones de seguridad de LLM a gran escala cuando el índice de atribución se reutiliza en muchas consultas, sin requerir los costos prohibitivos de almacenamiento o E/S de los métodos anteriores.

El artículo reconoce limitaciones, señalando que la construcción del índice de atribución aún requiere un paso de gradiente sobre el corpus (costo de preprocesamiento) y que los experimentos actuales a gran escala se centran en datos de ajuste fino supervisado en lugar de en los corpus completos de preentrenamiento. Sin embargo, el trabajo establece una nueva línea base para la atribución escalable, permitiendo el análisis de modelos que anteriormente se consideraban fuera de alcance para tales métodos.