Source-linked AI summary
TEE-X: TEE-aware Acceleration Framework for Large Vision Models at the Edge
Kurt M Wilson, Mohaiminul Al Nahian, Abeer Matar A. Almalky, Sadat Shahriyar, Souvik Kundu, Zhishan Guo, Abdullah Al Arafat, Adnan Siraj Rakin
TL;DR
Large vision models need TEE protection for confidentiality and integrity, but secure deployment is constrained by enclave memory and inference overhead. TEE-X keeps full ViT inference inside trusted memory using sensitivity-aware modularization, budget allocation, and SIMD vectorization, achieving near-GPU latency with high accuracy. Its performance advantage is most applicable to single-batch inference, while larger batches widen the GPU gap.
Problem
Large vision models require TEE-based confidentiality and integrity, but limited secure memory and performance overhead complicate fully protected, time-sensitive edge inference.
Method
TEE-X uses sensitivity-aware modularization, TEE-aware computation-budget allocation, and Arm TrustZone SIMD vectorization to execute complete ViT inference inside secure memory.
Results
TEE-X achieves near-GPU inference latency with high accuracy while keeping the entire model execution within trusted memory.
Takeaways & Limitations
TEE-X provides a practical approach to secure edge inference that avoids the security limitations of partitioning- or offloading-based deployment.
Takeaways & Limitations
For batch sizes beyond 1, the performance gap increases because GPU parallelism spans multiple inputs while TEE-X parallelism is column-wise within one input.
Abstract
from arXiv · showhide
Despite their remarkable success, machine learning models, particularly in vision applications, are alarmingly vulnerable to a range of security threats. One key factor in the attack landscape is the distinction between white-box and black-box threat models, as the latter poses challenges that limit attack effectiveness when access to model information is limited. As a result, using Trusted Execution Environments (TEEs) enhances security for machine learning applications by protecting model confidentiality and execution integrity, effectively shifting the execution environment from the white-box to the black-box side of the threat model spectrum. While adopting TEEs for large vision models, e.g., Vision Transformers (ViTs), is crucial for enhancing security and privacy, significant challenges related to memory constraints and increased computational latency must be addressed, especially in time-sensitive edge applications where safety and privacy are paramount. The objective of this work is to enable large vision models to be fully hosted within TEEs, achieving GPU-level inference latency for time-sensitive edge vision applications while maintaining performance. To this end, we propose TEE-X, a TEE-aware acceleration framework that introduces a sensitivity-aware modularization technique and enables vectorization in TEE inference. This design is validated on OP-TEE for Arm TrustZone, configured to optimize performance on the NVIDIA Jetson AGX Xavier for efficient edge vision applications using ViT models. The findings reveal that TEE-X delivers an effective TEE-aware acceleration framework that achieves minimal accuracy-latency trade-offs while ensuring fast and secure edge inference for vision models.
1 Introduction
Vision models face broad security threats, while TEEs improve confidentiality and integrity by moving execution toward a black-box threat model. TEE-X targets the memory and latency challenges of fully hosting large vision models in TEEs through sensitivity-aware modularization and vectorized inference.
- Vision models remain vulnerable to memory fault injections, sidechannels, and adversarial attacks.
- White-box attacks assume access to model architecture and parameters, whereas black-box attacks restrict adversaries to direct or indirect observations.
- TEEs protect model confidentiality and inference integrity by isolating execution and shifting the attack surface toward the black-box spectrum.
- Prior TEE deployments reduce memory pressure through model partitioning or selective operation placement, but do not necessarily keep the entire model inside the enclave.
- TEE-X combines sensitivity-aware modularization, TEE-aware computation-budget assignment, and Arm Neon SIMD vectorization for fully protected ViT inference.
2 Background and Related Work
TEEs securely isolate sensitive data, models, and computations but impose limited memory capacity and performance overhead. Prior vision-model frameworks often overlook ViT-specific challenges or execute only part of the model inside the enclave.
- TEEs are hardware-isolated environments that protect sensitive data, models, and computations from external access.
- TEE deployments face limited memory capacity and performance overhead when hosting large machine-learning models.
- The paper assumes a trusted enclave within a potentially untrusted host system, while excluding side-channel attacks from its threat model.
- Existing vision-model frameworks largely overlook ViT-specific architectural and system challenges.
- Many prior approaches execute only some model components inside the enclave instead of fully encapsulating the model.
3 TEE-X: TEE-aware Acceleration Framework
TEE-X enables full ViT inference inside TEEs by combining sensitivity-aware modularization, budget allocation, learned weight transformations, and vectorized computation. Its theoretical budget can match GPU-level performance for some single-batch configurations, while larger batches widen the performance gap.
- TEE-X combines sensitivity-aware modularization, TEE-aware computation budgets, learned weight transformations, and vectorization to accelerate complete in-enclave ViT inference.The framework stores transformed weights and module-wise linear transforms during inference, reducing TEE memory and computation requirements.
- Modularization: Layers with similar sensitivity are grouped into modules, then assigned computation budgets so sensitive layers retain accuracy while less sensitive layers reduce latency.Sensitivity is estimated by perturbing one layer at a time and measuring the resulting accuracy drop.
- Modularization: Each weight matrix is transformed into Ai and Bi factors, with hi controlling the computation required for xi × wi = (xi × Ai) × Bi.The learned transformations decompose weight matrices under the available TEE computation budget.
- Transformation-network training: The transformed model is jointly trained with cross-entropy, knowledge-distillation, and mean-squared-error losses to preserve task performance and original layer outputs.Distillation preserves the pretrained ViT’s output behavior, while MSE aligns transformed and original layer outputs.
- Performance: TEE-X’s theoretical computation budget matches GPU-level performance in practical TEE implementations for some single-batch configurations.For batch sizes beyond one, the performance gap increases because GPUs parallelize multiple inputs, whereas TEE-X parallelizes columns within a single input.
4 Experimental Evaluation
TEE-X is evaluated as a complete in-enclave ViT inference framework on OP-TEE and Jetson AGX Xavier, targeting secure execution with practical memory, accuracy, and latency trade-offs. Across datasets, comparisons, and ablations, sensitivity-aware modularization and Neon vectorization provide efficient acceleration while preserving model performance and reducing exposure to attacks.
- Evaluation Setting: TEE-X uses OP-TEE on Arm TrustZone with an OP-TEE host application and trusted application evaluated on an Nvidia Jetson AGX Xavier.Experiments use DeiT-Small and DeiT-Base on CIFAR-10 and ImageNet.
- Accuracy and Latency Trade-offs: TEE-X achieves substantial model-size reduction while keeping accuracy degradation below 1% at C = 0.5 and C = 0.4.At C = 0.2, DeiT-Small is compressed approximately 47× with a 2.65% accuracy drop.
- Accuracy and Latency Trade-offs: TEE-X provides a flexible accuracy–latency trade-off: C = 0.5 preserves near-baseline accuracy with comparable GPU latency, while C = 0.2 can reduce latency further and outperform GPU execution.The reported trade-off relies on modularization and vectorization support.
- Comparison with Competitive Methods: TEE-X achieves comparable or lower latency than GPU execution while executing the full model inside the TEE, unlike partitioned or operation-split frameworks.Prior approaches incur latency gaps from enclave transitions and communication overhead.
- Comparison with Competitive Methods: On ImageNet, TEE-X reduces DeiT-Small to 4.63 MB with 76.44 accuracy, outperforming the listed pruning and 2-bit quantization results.The reduced DeiT-Small model remains competitive at 3.81 MB on ImageNet, and results extend across datasets and architectures.
- Security Evaluation: Black-box attack success rates are considerably lower than white-box rates, indicating reduced model exposure under TEE-X deployment.The evaluation reports this transition toward a black-box attack spectrum as a security benefit.
- Ablation Study: TEE-X achieves 1.08s inference latency, while removing modularization increases latency to 1.8s and removing vectorization increases it to 1.77s.Accuracy remains unchanged when modularization is removed, isolating its latency contribution.
- Ablation Study: The proposed computation-budget assignment preserves model performance better than uniform allocation, especially under strict budgets.The strategy allocates more budget to sensitive layers and reduces budget in less sensitive regions.
5 Conclusion
TEE-X enables large vision models to run entirely within trusted memory while targeting near-GPU inference latency and high accuracy for edge applications.
- TEE-X enables efficient inference of large vision models within secure memory for edge applications.