Inferencia de aprendizaje automático cifrada de extremo a extremo con Amazon SageMaker AI y FHE

La inferencia del aprendizaje automático (ML) a menudo requiere el procesamiento de datos confidenciales: registros médicos, información comercial patentada o comunicaciones personales. ¿Qué pasaría si pudiera ejecutar la inferencia de ML en la nube mientras oculta sus datos de la propia nube? Más específicamente, ¿qué pasaría si pudiera exigir que sus datos permanecieran cifrados durante todo el proceso de inferencia de ML? Esta publicación le mostrará cómo utilizar Amazon SageMaker AI con cifrado totalmente homomórfico (FHE) para realizar inferencia de ML. Usando FHE, presentamos un enfoque para la inferencia de ML que está diseñado para mantener las consultas, respuestas y valores intermedios cifrados e ilegibles para los observadores, incluido el propio SageMaker AI.

FHE es una forma de cifrado que permite que los datos cifrados se procesen de forma cifrada sin descifrarlos. En la configuración de inferencia de ML, puede usarlo para aplicar un modelo a una consulta cifrada sin descifrarla, lo que produce una predicción cifrada. Considere estos escenarios en los que dicha capacidad proporcionaría valor:

Atención médica: una compañía de seguros de salud desea proporcionar a los médicos un modelo de aprendizaje automático que prediga los resultados de los procedimientos médicos basándose en datos de diagnóstico. La publicación del modelo en la nube simplifica la implementación, pero los médicos no pueden exponer la información médica del paciente a terceros debido a las normas de privacidad. Sector energético: una corporación de petróleo y gas utiliza ML para evaluar fotografías satelitales de posibles sitios de perforación y seleccionar fotografías para una evaluación adicional de expertos. Quieren alojar el modelo en la nube para ahorrar costos, pero no pueden exponer fotografías de lugares políticamente sensibles a terceros. Telecomunicaciones: un operador de telecomunicaciones quiere procesar los correos electrónicos de los clientes para detectar spam y phishing. Necesitan ML basado en la nube para lograr escalabilidad, pero las regulaciones de protección de datos exigen que los mensajes de los clientes permanezcan cifrados en terceros.

Este blog analizó anteriormente FHE para la inferencia de ML en la publicación Habilite el cifrado totalmente homomórfico con puntos finales de Amazon SageMaker para una inferencia segura en tiempo real, pero esta publicación va un poco más allá. Esa publicación anterior mostró cómo implementar la inferencia basada en FHE 'desde cero' mediante la elaboración manual de un algoritmo de regresión lineal utilizando una biblioteca de bajo nivel llamada SEAL. En cambio, esta publicación muestra un enfoque mucho más flexible y de mayor nivel basado en concrete-ml, una biblioteca de alto nivel creada específicamente para la inferencia basada en FHE. Admite varios tipos comunes de modelos listos para usar e incluso es compatible con API con la conocida biblioteca ML scikit-learn.

En esta publicación, aprenderá cómo:

Entrene un modelo de ml concreto en SageMaker AI usando un contenedor personalizado Implemente ese modelo en un punto final de inferencia de SageMaker AI Cree un cliente personalizado para la inferencia de ml concreto Utilice ese cliente para realizar consultas a su punto final de inferencia

Cuando termine, tendrá un sistema que utiliza concrete-ml en SageMaker AI diseñado para realizar inferencia de ML cifrada de un extremo a otro.

Descripción general de la solución

El uso de concrete-ml en SageMaker AI funciona de la siguiente manera:

El propietario del modelo prepara sus datos para el entrenamiento. Concrete-ml funciona bien cuando todas las funciones se han normalizado a la misma escala, como [-1, 1]. El propietario del modelo utiliza estos datos para entrenar una versión de su modelo habilitada para FHE. Este modelo está diseñado para realizar cálculos sobre datos cifrados en lugar de texto sin formato. El propietario del modelo aloja este modelo en SageMaker AI. Los clientes cifran sus consultas utilizando el esquema FHE admitido por el modelo. Los clientes envían consultas cifradas al modelo habilitado para FHE en la nube. El modelo transforma la consulta cifrada en una predicción cifrada sin descifrar los valores durante el cálculo FHE. El modelo devuelve la respuesta cifrada al cliente, quien la descifra para recuperar la predicción.

Esto difiere y complementa los entornos informáticos confidenciales como los proporcionados por el sistema Nitro de Amazon Web Services (AWS) en Amazon Elastic Compute Cloud (Amazon EC2). Con AWS Nitro Enclaves, las consultas se descifran y procesan en texto sin formato dentro de entornos aislados y reforzados que brindan aislamiento de CPU y memoria. Con FHE, las consultas permanecen cifradas en todo momento; La seguridad se basa en las matemáticas más que en el hardware o el software.

Requisitos previos

Para implementar esta solución, necesita:

Un entorno de desarrollo local con Python 3.12 instalado, la capacidad de instalar paquetes usando pip y Docker u otro software de creación de contenedores instalado localmente. Además, estas instrucciones te recomendarán trabajar en entornos virtuales, pero esto no es estrictamente necesario. Una cuenta de AWS que contiene:

Le sugerimos seguir las mejores prácticas de seguridad para Amazon S3.

Roles en AWS Identity and Access Management (IAM) para El creador del modelo El creador del punto final de inferencia El propio punto final de inferencia Los clientes

Encuentre políticas de IAM para estos roles, junto con un ejemplo trabajado para el corpus MNIST de dígitos escritos a mano, en el repositorio de código de muestra.

Antes de comenzar, tenga en cuenta que en el momento de escribir este artículo, concrete-ml está disponible en Zama para creación de prototipos o uso no comercial sin necesidad de una licencia paga. Sin embargo, es posible que necesite una licencia comercial para uso comercial.

Capacitación

Construya e implemente el contenedor de capacitación

Para construir el contenedor de entrenamiento:

Asume el rol de formador de modelos. Cree un archivo Dockerfile.training localmente. Agregue el siguiente contenido a Dockerfile.training:

Verifique que los números de versión coincidan en todo el sistema. La biblioteca concrete-ml requiere paridad de versiones en todo el sistema para Python, el paquete concrete-ml y el paquete concrete-python.

Construya la imagen del contenedor:

Envíe la imagen a Amazon ECR: ejecute el comando de autenticación para iniciar sesión en Docker en su registro de Amazon ECR:

Etiqueta la imagen con el nombre de tu repositorio:

Empuje la imagen etiquetada:

Verificar que el contenedor esté disponible.

Debería ver la salida JSON que contiene su imagen con un campo imageDigest no vacío y la última etiqueta.

Entrenar el modelo

Para entrenar el modelo, complete lo siguiente.

Nota: en estos pasos, concrete-ml no es diferente de cualquier otro marco de ML y el contenedor de capacitación no es diferente de cualquier otro contenedor de capacitación personalizado. Tenga en cuenta que el entrenamiento se produce sobre datos de texto sin formato. Es decir, concrete-ml no requiere un procesamiento previo de estos datos más allá de la normalización. Pero si es necesario un preprocesamiento adicional para la capacitación regular, sigue siendo necesario aquí (y debe ocurrir antes o como parte del trabajo de capacitación).

Crear el guión de entrenamiento

Cree un archivo llamado Training_script.py. Agregue el siguiente código de plantilla a Training_script.py:

Implemente la lógica de carga de datos en la función do_training. Implemente la lógica de entrenamiento del modelo en la función do_training.

Crear un marco personalizado

Para mayor comodidad, le recomendamos que cree un marco personalizado para integrar su contenedor de capacitación en SageMaker AI. Para hacerlo:

Cree un archivo llamado framework.py. Agregue el siguiente contenido a framework.py:

Actualice el valor image_uri con la ubicación de su contenedor de capacitación de Amazon ECR.

Iniciar el trabajo de formación

Esta sección mostrará cómo iniciar el trabajo de capacitación con un script de Python, pero también se puede hacer usando la consola o la interfaz de línea de comandos de AWS (AWS CLI). (Nota: los trabajos de capacitación generan cargos según el tipo de instancia y la duración).

Cree un entorno virtual para Python 3.12. Activa el entorno virtual. Instale los siguientes paquetes usando pip:

Cree un archivo llamado start_training.py. Agregue el siguiente contenido a start_training.py:

Actualice los valores de tipo_instancia, rol, ruta_salida, ubicación_código y entradas con su configuración específica. Ejecute este archivo: verifique que la capacitación se haya completado exitosamente verificando el estado del trabajo de capacitación:

Busque TrainingJobStatus: Completado. Luego verifique que los archivos de salida existan:

Confirme que server.zip y client.zip estén presentes.

Una vez completada la capacitación, el contenedor de capacitación guarda dos archivos en el depósito del modelo: server.zip (usado por el punto final de inferencia) y client.zip (usado por los clientes para cifrar consultas).

Inferencia

Construya e implemente el contenedor de inferencia

La inferencia de ML basada en FHE será más compleja que la inferencia de ML estándar debido a algunas limitaciones técnicas nuevas:

Los clientes necesitan información específica del modelo de client.zip para generar claves criptográficas. Los textos cifrados FHE pueden exceder los límites de tamaño de las consultas de SageMaker AI, por lo que el cliente y el servicio deben comunicarlos fuera de las llamadas a la API de SageMaker AI. La evaluación de FHE puede tardar más que los tiempos de espera de SageMaker AI, por lo que la inferencia utilizará los mecanismos de SageMaker AI para la inferencia asincrónica. El punto final necesita una clave de evaluación (un tipo de clave pública) del cliente para realizar la evaluación FHE.

Para adaptarse a estos nuevos requisitos y optimizar la experiencia del usuario, le mostramos cómo construir un sistema en el que

Un cliente personalizado cifra las consultas y les adjunta claves de evaluación. Un punto final de entrenamiento personalizado recupera client.zip cuando es necesario y lo utiliza para evaluar el modelo FHE. El mismo cliente personalizado descifra las predicciones del punto final de entrenamiento. El cliente y el punto final se comunican textos cifrados y claves entre sí mediante Amazon S3.

Para implementar y utilizar este sistema, complete las siguientes secciones.

Escribe tu predictor

Cree un archivo llamado predictor.py con el siguiente contenido.

Este predictor espera que la 'consulta' contenga tres ubicaciones de Amazon S3: dos para encontrar la consulta cifrada y la clave de evaluación asociada, y una para escribir la predicción. Descarga la consulta y la clave, evalúa el modelo FHE en ellas y escribe la predicción en Amazon S3.

Empaquetar el predictor en un contenedor

Para empaquetar este predictor en un contenedor:

Asuma el rol de creador de endpoints. Cree un nuevo directorio para los archivos contenedores. Copie predictor.py en el nuevo directorio. Obtenga los archivos repetitivos necesarios (nginx.conf, server y wsgi.py) descargándolos del repositorio de muestra o copiándolos de la documentación de SageMaker AI para contenedores de inferencia personalizados. (Nota: en este último caso, aumente el valor de tiempo de espera en nginx.conf para permitir que se complete la evaluación FHE). Cree una inferencia Dockerfile.en ese directorio. Agregue el siguiente contenido al archivo Dockerfile.inference:

Construya la imagen del contenedor:

Envíe la imagen a Amazon ECR. Ejecute el comando de autenticación para iniciar sesión en Docker en su registro de Amazon ECR:

Etiqueta la imagen con el nombre de tu repositorio:

Empuje la imagen etiquetada:

Verifique que el contenedor esté disponible:

Debería ver la salida JSON que contiene su imagen con un campo imageDigest no vacío y la última etiqueta.

Implementar el punto final de inferencia

(Importante: los puntos finales incurren en cargos continuos hasta que se eliminan, y los costos variarán según el tipo de instancia, la duración del entrenamiento y el tiempo de actividad del punto final. Para obtener información detallada sobre precios, consulte Precios de IA de Amazon SageMaker. Recuerde eliminar el punto final cuando finalice para evitar costos innecesarios). Continuar usando el rol de creador de puntos finales:

Crea un entorno virtual. Activa este entorno virtual. Utilice pip para instalar los siguientes paquetes:

Cree un archivo start_inference_endpoint.py con el siguiente contenido:

Ejecute el script:

Verifique que el punto final esté en servicio:

Espere hasta que EndpointStatus muestre InService antes de continuar. Esto puede tardar varios minutos.

El script imprimirá el nombre del punto final. Registre este nombre para el cliente.

crear el cliente

El usuario no debería necesitar saber nada sobre FHE para utilizar su sistema. Por tanto, el cliente ocultará todos los datos de FHE. En concreto, el cliente:

Recupere client.zip de Amazon S3. Utilice client.zip para generar claves. Cifre la consulta con esas claves. Escriba la consulta cifrada y la clave de evaluación asociada en Amazon S3. Envíe estas ubicaciones al punto final de inferencia y reciba la ubicación de Amazon S3 de la predicción cifrada. Recupere la predicción cifrada y descifrela.

Para crear este cliente:

Cree un archivo llamado client.py. Agregue el siguiente código de plantilla a client.py:

Implemente la función get_query() para recuperar su consulta en texto plano. Actualice los valores de los marcadores de posición para las ubicaciones de Amazon S3, el nombre del punto final y la ubicación del modelo. Agregue código de manejo de excepciones para los bloques de marcador de posición para administrar TimeoutError, FileNotFoundError y TarError de acuerdo con los requisitos de su aplicación.

(Es posible que haya notado que el cliente y el punto final tratan las consultas y respuestas cifradas de manera diferente. Los clientes envían consultas cifradas a los puntos finales escribiéndolas manualmente en Amazon S3 y enviando la ubicación de Amazon S3 como la consulta real. Los puntos finales envían resultados cifrados directamente, lo que permite a SageMaker AI manejar la escritura/lectura desde Amazon S3. ¿Por qué la diferencia? La respuesta cifrada es una sola cadena de bytes, que SageMaker AI puede manejar de forma natural. La consulta cifrada, sin embargo, es una estructura JSON que debe contener la ubicación de las claves de evaluación. La consulta cifrada debería codificarse (como con Base64) para estar incrustada en el mismo JSON, lo que agrega tiempo de red y procesamiento innecesarios. Por lo tanto, el código de muestra omite este paso de codificación al manejar las consultas cifradas por sí mismo.

Entonces:

Crea un entorno virtual. Activa el entorno virtual. Instale los paquetes necesarios:

Finalmente:

Asume el rol de cliente. Ejecute este script:python client.py Verifique que el cifrado FHE esté funcionando correctamente comparando el resultado de la predicción con los resultados esperados.

Limpiar recursos

Para evitar incurrir en cargos futuros, elimine los recursos que creó:

Elimine el punto final de inferencia a través de la consola o SDK de SageMaker AI. Verifique que el punto final haya sido eliminado:

Esto debería devolver un error que indica que el punto final no existe.

Elimine la configuración del punto final a través de la consola o SDK de SageMaker AI. Verifique que la configuración del punto final se haya eliminado:

Esto no debería mostrar ninguna configuración de punto final coincidente.

Elimine el modelo de SageMaker AI a través de la consola o SDK de SageMaker AI. Verifique que el modelo haya sido eliminado:

Esto no debería mostrar ningún modelo coincidente.

Elimine los artefactos del modelo, las consultas cifradas, las respuestas cifradas y las claves de evaluación de Amazon S3 a través de la consola de Amazon S3 o la AWS CLI. Verifique que los objetos de Amazon S3 se hayan eliminado:

Esto debería mostrar objetos vacíos o sin coincidencias.

Elimine las imágenes del contenedor de Amazon ECR a través de la consola de Amazon ECR o la CLI de AWS. Verifique que las imágenes del contenedor fueron eliminadas:

Esto no debería mostrar imágenes coincidentes.

Problemas comunes

TimeoutError durante la inferencia: aumente WaiterConfig max_attempts o utilice un tipo de instancia más grande. Errores de acceso denegado: verifique que los roles de IAM tengan los permisos correctos de S3 y SageMaker AI. Errores de compilación del contenedor: verifique que Docker tenga suficiente memoria (más de 8 GB). Errores del servidor durante la inferencia: verifique la paridad de versiones entre paquetes concrete-ml.

Consideraciones de rendimiento y seguridad

FHE proporciona protección criptográfica pero conlleva compensaciones en el rendimiento. La sobrecarga depende del modelo, pero normalmente se pueden esperar ralentizaciones de hasta 100.000 veces en comparación con la inferencia de texto sin formato. Puede reducir esta desaceleración de varias maneras. La primera es aumentar la cantidad de vCPU en la instancia. Otra es utilizar una técnica de ML estándar llamada "cuantización" que reduce la precisión numérica utilizada en la inferencia de modelos. Debido a que el tiempo de ejecución de concrete-ml aumenta con la precisión numérica, la cuantificación podría ayudar al rendimiento aquí incluso más que en la inferencia ML normal. La cuantificación puede reducir la precisión del modelo, que de otro modo no se ve afectada por la conversión a FHE. Sin embargo, la cuantificación en el código del modelo redujo la sobrecarga a 2800X (67 ms a 187 s en una instancia ml.m5.xlarge) sin pérdida observable de precisión. Al aumentar la cantidad de vCPU, puede reducirla aún más a 500 veces (46 en una instancia ml.m5.24xlarge).

Esto sigue siendo una desaceleración significativa para algunas aplicaciones. Debido a esta sobrecarga, FHE aún no es adecuado para aplicaciones interactivas sensibles a la latencia. Sin embargo, puede resultar práctico para cargas de trabajo de procesamiento por lotes o asíncronas donde los requisitos de privacidad superan las preocupaciones de latencia. Por ejemplo, considere los casos de uso del comienzo de esta publicación:

Proporcionar a los médicos un modelo de aprendizaje automático que predice los resultados de los procedimientos médicos basándose en datos de diagnóstico. Evaluación de fotografías satelitales de posibles sitios de perforación de petróleo/gas para seleccionar fotografías para una evaluación adicional de expertos. Detección de spam y phishing en mensajes de correo electrónico.

Cada uno de estos casos de uso puede tolerar unos segundos adicionales de latencia.

Es importante que los clientes mantengan en secreto las consultas y predicciones descifradas, ya que un cifrado de ml concreto y su descifrado de texto plano (cuando se combinan) podrían revelar información sobre la clave de cifrado secreta. Además, es importante saber que este sistema no protege el secreto del modelo. Las consultas y respuestas estarán cifradas y opacas para SageMaker AI, pero concrete-ml no cifra el modelo en sí. Es posible que el modelo aún sea visible para Sagemaker AI. También podría ser susceptible a ataques de "robo de modelos" por parte de quienes pueden ver consultas y respuestas en texto plano. Por último, concrete-ml no proporciona privacidad del circuito: es posible que los cipertextos puedan revelar información sobre el modelo. Sin embargo, los clientes aún pueden proteger los modelos y los textos cifrados con los mecanismos de seguridad estándar que AWS proporciona para Amazon S3 y SageMaker AI. Recuerde: la seguridad es una responsabilidad compartida entre AWS y cada cliente. De acuerdo con las mejores prácticas, los clientes deben:

Siga el principio de privilegio mínimo al crear roles de IAM. Otorgue solo los permisos mínimos necesarios para que cada rol realice su función. Revise las políticas de IAM de muestra en el repositorio y ajuste los ARN de recursos y las acciones para que coincidan con su caso de uso específico. Habilite el cifrado de depósitos de Amazon S3 para valores que no sean textos cifrados FHE. Esto incluye habilitar el cifrado predeterminado en todos los depósitos de Amazon S3 que almacenan modelos, datos y claves de evaluación para proteger los datos en reposo. Reduzca los permisos del depósito de Amazon S3 al mínimo requerido por el sistema.

Conclusión

Puede utilizar herramientas basadas en FHE en SageMaker AI para realizar inferencias sobre datos cifrados diseñados para permanecer ilegibles durante todo el proceso. Este enfoque puede brindarle los beneficios de SageMaker AI (agilidad, escalabilidad e infraestructura administrada) mientras lo ayuda a mantener la protección criptográfica desde la consulta hasta la respuesta.

Para obtener más información sobre la seguridad y el cifrado en AWS, consulte los siguientes recursos:

Si tiene preguntas o comentarios, contáctenos en aws-crypto-compute@amazon.com.

Sobre los autores