AttestLLM: Efficient Attestation Framework for Billion-scale On-device LLMs

1. El Problema

La implementación de Modelos de Lenguaje Grande (LLMs) en dispositivos locales (on-device) ofrece ventajas significativas en privacidad y latencia. Sin embargo, esto introduce un desafío de seguridad crítico: la protección de la Propiedad Intelectual (IP) a nivel de hardware.

• Amenaza: Terceros no confiables en la cadena de suministro (como revendedores) pueden acceder al dispositivo y reemplazar los LLMs optimizados por el fabricante con modelos no autorizados, ilegales o inseguros.
• Limitaciones de soluciones anteriores: Las técnicas de atestación (verificación de legitimidad) existentes están diseñadas para redes neuronales pequeñas (millones de parámetros) y no escalan a LLMs de miles de millones de parámetros debido a:
  1. Restricciones de hardware: Los entornos de ejecución confiable (TEE) en dispositivos móviles (ej. Arm TrustZone) tienen memoria segura muy limitada (10-32 MB), insuficiente para cargar los parámetros y activaciones de un LLM completo.
  2. Eficiencia: La verificación de marcas de agua (watermarks) debe ser computacionalmente ligera para no degradar el rendimiento de inferencia, pero lo suficientemente robusta para resistir ataques.

2. Metodología: AttestLLM

AttestLLM es el primer marco de atestación diseñado específicamente para proteger la IP de hardware de LLMs a escala de miles de millones de parámetros. Utiliza un enfoque de co-diseño de algoritmo, software y hardware dividido en dos fases principales:

A. Marcado de Agua Offline (Watermarking)

El objetivo es incrustar firmas específicas del dispositivo en el modelo sin degradar su calidad.

• Análisis de Sensibilidad: Se evalúa cada capa del LLM para determinar su sensibilidad. Las capas con activaciones más grandes (más informativas) reciben menos bits de la marca de agua, mientras que las capas menos sensibles reciben más. Esto se calcula mediante la magnitud de las activaciones pico.
• Inserción Optimizada (Pre y Post Cuantización):
  • Se utiliza un algoritmo de optimización basado en la cuantización (necesaria para dispositivos móviles).
  • Fase Pre-cuantización: Se ajusta el modelo de precisión completa usando descenso de gradiente para aproximar la firma objetivo.
  • Fase Post-cuantización: Dado que los pesos se vuelven discretos tras la cuantización, se emplea optimización de orden cero (zeroth-order) para refinar la inserción de la firma sin romper la fidelidad del modelo.
• Verificación Independiente: Cada capa del LLM tiene su propia marca de agua. Durante la inferencia, se guardan las activaciones de entrada y se verifica la firma en la salida de cada capa individualmente.

B. Atestación Online en TEE

La verificación ocurre dentro del Entorno de Ejecución Confiable (TEE) del dispositivo.

• Ampliación de Memoria Segura: Utiliza una enclave basada en virtualización (pKVM) para superar las limitaciones de memoria de TrustZone, permitiendo manejar LLMs grandes.
• Muestreo Dinámico y Paralelo:
  • No se verifica todo el modelo en cada ronda. El sistema selecciona aleatoriamente un subconjunto de capas ($k$) de las capas totales ($|M|$).
  • Estas capas se copian de manera segura desde el entorno de ejecución rico (REE) al TEE.
  • Se utiliza un enfoque pipelined y paralelo: mientras un hilo realiza la copia segura y descifrado de claves, otros hilos verifican las marcas de agua de las capas previamente copiadas.
• Mecanismos de Seguridad:
  • Aleatoriedad: La selección de capas a verificar se genera aleatoriamente dentro del TEE para evitar que un atacante prevea qué capas serán revisadas.
  • Monitoreo del Hipervisor: Detecta modificaciones en las páginas de memoria asignadas al modelo en tiempo de ejecución.
  • Aleatorización de Memoria: Las memorias de los pesos se mezclan (shuffled) para dificultar la búsqueda por fuerza bruta de las marcas de agua.

3. Contribuciones Clave

1. Primero en su clase: Es el primer marco de atestación capaz de proteger la IP de hardware de LLMs de escala masiva (miles de millones de parámetros) en dispositivos móviles.
2. Co-diseño Escalable: Combina técnicas de marcaje de agua robustas y preservadoras de calidad con optimizaciones de sistema (virtualización, paralelismo) para operar bajo restricciones estrictas de memoria y energía.
3. Alta Eficiencia y Robustez: Logra una verificación fiable con una sobrecarga mínima, resistiendo ataques de falsificación, reemplazo de modelos y manipulaciones del sistema TEE.

4. Resultados Experimentales

El marco se evaluó en varios LLMs on-device (Llama3, Qwen, Phi) cuantizados a INT4 e INT8 en un dispositivo ARM.

• Fiabilidad y Precisión:
  • Tasa de Extracción de Marca de Agua (WER): 100% en todos los modelos probados.
  • Tasa de Aprobación de Atestación (APR): 100% para modelos legítimos; 0% para modelos no marcados o manipulados.
• Preservación de Calidad (Fidelidad):
  • La inserción de la marca de agua tiene un impacto mínimo en el rendimiento del modelo.
  • Degradación promedio de Perplejidad (PPL): 0.62% (INT4) y 0.07% (INT8).
  • Degradación de Precisión Zero-shot: 0.09% (INT4) y 0.15% (INT8).
• Eficiencia (Sobrecarga):
  • Comparado con métodos anteriores de "infección TEE" (TEE-shield inference), AttestLLM reduce la sobrecarga de latencia en un 62x y la de energía en un 30x.
  • Sobrecarga de latencia absoluta: ~9.3% (INT4) y ~8.3% (INT8).
• Resistencia a Ataques:
  • Falsificación: Resistente porque las claves y el conjunto de datos de activación (trigger dataset) están cifrados y ocultos.
  • Reemplazo de Modelo: Detecta incluso si el atacante modifica solo un subconjunto de capas (ej. 4 de 32), reduciendo la probabilidad de evasión a niveles insignificantes ($< 10^{-3}$) tras varias rondas de verificación.
  • Ataques al Sistema TEE: Mitigados mediante temporizadores confiables y monitoreo del hipervisor que impiden la ejecución si no hay interacción de atestación.

5. Significado

AttestLLM representa un avance crucial para la adopción segura de IA en el borde (edge AI). Permite a los fabricantes de hardware y proveedores de modelos distribuir LLMs optimizados con la garantía de que no serán reemplazados ni manipulados por actores maliciosos en la cadena de suministro. Al resolver el dilema entre la seguridad robusta y la eficiencia en recursos limitados, habilita un ecosistema de confianza donde la propiedad intelectual de hardware y software puede protegerse sin sacrificar la experiencia del usuario ni el rendimiento del dispositivo.