Efficient Implementation of AES on AI Accelerators for Machine Learning Model Protection

Sammanfattning

The widespread deployment of machine learning models on end-user devices raises important concerns regarding the protection of sensitive data and model intellectual property. Training such models requires significant computational resources, large datasets, and careful architectural design, making trained models valuable assets vulnerable to extraction and unauthorized reuse. On Apple platforms, machine learning inference often relies on hardware accelerators such as the GPU and the Apple Neural Engine (ANE), creating new opportunities for on-device cryptographic protection. This master thesis explores the design and implementation of cryptographic methods aimed at protecting machine learning models by performing cryptographic computations directly on Apple Silicon accelerators. The primary focus is on implementing the Advanced Encryption Standard (AES) on Apple GPUs and on the ANE, with the objective of enabling secure model deployment through encrypted weights. In this thesis, we develop the first known AES implementations specifically targeting Apple GPUs using the Metal framework, based on different algorithmic strategies. These implementations are benchmarked against CPU-based solutions, showing that Apple GPUs can reach very high encryption throughput and achieve performance comparable to state-of-the-art GPU-based implementations reported in the literature. In addition, this work presents the first known implementations of AES running entirely on the ANE. Multiple ANE-based designs are proposed and analyzed, demonstrating that cryptographic computations can be performed exclusively on the ANE without CPU involvement. Building on these results, this thesis introduces a complete model hardening method for ANE-based machine learning models. In this approach, model weights are stored in encrypted form and decrypted dynamically at inference time directly on the ANE, with weights decryption as part of the model architecture. This allows models to be distributed securely while keeping weights protected during execution. Experiments show that this approach is feasible for small models, while also highlighting the challenges of scaling to larger architectures.