Securing Transformer-based AI Execution via Unified TEEs and Crypto-protected Accelerators

📄 arXiv: 2507.03278v2 📥 PDF

作者: Jiaqi Xue, Yifei Zhao, Mengxin Zheng, Fan Yao, Yan Solihin, Qian Lou

分类: cs.CR, cs.LG

发布日期: 2025-07-04 (更新: 2025-07-13)

备注: 15 pages


💡 一句话要点

TwinShield:通过统一TEE和加密加速器保护Transformer模型安全执行

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: Transformer模型 可信执行环境 GPU加速 安全推理 同态加密 机器学习即服务 数据安全 模型保护

📋 核心要点

  1. 现有基于TEE的Transformer模型推理方案,在保护数据和模型安全的同时,面临着性能瓶颈,尤其是在大型语言模型中,计算和内存需求巨大。
  2. TwinShield框架通过安全地将Transformer模型中的关键操作(如Attention和SoftMax)卸载到GPU加速器,实现了数据和模型的双重保护。
  3. 实验结果表明,TwinShield能够将约87%的计算卸载到GPU,并在各种Transformer模型上实现了4.0x - 6.1x的速度提升,显著优于之前的方案。

📝 摘要(中文)

Transformer模型,如大型语言模型(LLMs),在人工智能任务中取得了巨大突破,并在安全关键领域得到广泛应用。由于其规模庞大和开发成本高昂,这些模型成为AI利益相关者极具价值的知识产权,并越来越多地通过机器学习即服务(MLaaS)进行部署。然而,MLaaS通常在不受信任的云基础设施上运行,使数据和模型面临潜在的泄露风险。主流的保护机制利用可信执行环境(TEEs),通过硬件加密和完整性检查来保护机密数据的机密性和完整性。不幸的是,完全在TEE中运行模型推理会造成显著的减速,这在LLM中由于涉及大量的计算和内存占用而进一步加剧。最近的研究表明,将部分模型推理操作卸载到不受信任的加速器(如GPU)的混合TEE方案是一种有前途的解决方案。然而,先前的卸载方案未能确保Transformer推理中数据和模型的双重保护,因为它们无法安全地卸载关键操作,即Attention和SoftMax,迫使这些计算仍然局限于TEE中。为了应对这些挑战,我们提出了TwinShield,一个支持在异构TEE和加速器系统中安全Transformer推理的框架,为模型和数据提供双重保护。TwinShield将约87%的计算卸载到GPU,并在各种Transformer模型上实现了比以前的方法快4.0x - 6.1x的速度提升。

🔬 方法详解

问题定义:论文旨在解决在MLaaS环境下,Transformer模型(特别是LLM)在不受信任的云基础设施上运行时,数据和模型面临的安全风险问题。现有基于TEE的方案虽然能提供安全保障,但由于计算和内存开销巨大,导致性能显著下降。混合TEE方案虽然尝试将部分计算卸载到GPU,但无法安全地卸载Attention和SoftMax等关键操作,限制了性能提升。

核心思路:TwinShield的核心思路是通过设计一种新的框架,能够安全地将Transformer模型中的关键计算操作(包括Attention和SoftMax)卸载到不受信任的GPU加速器上,同时保证数据和模型的机密性和完整性。通过充分利用GPU的并行计算能力,缓解TEE内部的计算压力,从而提高整体推理性能。

技术框架:TwinShield框架包含以下主要模块:1) 安全的内存管理模块,用于在TEE和GPU之间安全地传输数据;2) 加密加速器接口,用于安全地执行Attention和SoftMax等关键操作;3) 任务调度器,用于优化TEE和GPU之间的任务分配。整体流程是,首先在TEE内部加载模型和数据,然后将部分计算任务卸载到GPU加速器,GPU执行完计算后,将结果返回给TEE,TEE进行后续处理。

关键创新:TwinShield最重要的技术创新点在于能够安全地将Attention和SoftMax等关键操作卸载到GPU加速器。之前的方案由于无法保证这些操作在GPU上的安全性,只能将其限制在TEE内部执行,严重限制了性能。TwinShield通过设计特殊的加密加速器接口和安全的数据传输机制,解决了这个问题。

关键设计:TwinShield的关键设计包括:1) 使用同态加密技术来保护在GPU上执行的Attention和SoftMax操作的数据;2) 设计了一种新的内存管理机制,确保TEE和GPU之间的数据传输是安全的,防止数据泄露;3) 优化了任务调度算法,以最大限度地利用GPU的计算资源。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TwinShield框架在各种Transformer模型上进行了实验评估,结果表明,TwinShield能够将约87%的计算卸载到GPU,并在各种Transformer模型上实现了4.0x - 6.1x的速度提升,显著优于之前的基于TEE的方案。此外,TwinShield在保证性能提升的同时,能够有效地保护数据和模型的安全性,防止数据泄露和模型篡改。

🎯 应用场景

TwinShield框架可应用于各种需要保护数据和模型安全性的Transformer模型推理场景,例如:金融风控、医疗诊断、智能客服等。通过在MLaaS平台上部署TwinShield,可以为用户提供安全、高效的Transformer模型推理服务,降低数据泄露的风险,并保护AI模型的知识产权。该研究的成果有助于推动Transformer模型在安全敏感领域的广泛应用。

📄 摘要(原文)

Recent advances in Transformer models, e.g., large language models (LLMs), have brought tremendous breakthroughs in various artificial intelligence (AI) tasks, leading to their wide applications in many security-critical domains. Due to their unprecedented scale and prohibitively high development cost, these models have become highly valuable intellectual property for AI stakeholders and are increasingly deployed via machine learning as a service (MLaaS). However, MLaaS often runs on untrusted cloud infrastructure, exposing data and models to potential breaches. Mainstream protection mechanisms leverage trusted execution environments (TEEs) where confidentiality and integrity for secretive data are shielded using hardware-based encryption and integrity checking. Unfortunately, running model inference entirely within TEEs is subject to non-trivial slowdown, which is further exacerbated in LLMs due to the substantial computation and memory footprint involved. Recent studies reveal that the hybrid TEE-based scheme offloading partial model inference operations to the untrusted accelerators (e.g., GPU) is a promising solution. However, prior offloading schemes fail to ensure dual protection of data and model in Transformer inference, as they cannot securely offload critical operations, i.e., Attention and SoftMax, forcing these computations to remain confined within TEEs. To address these challenges, we propose TwinShield, a framework enabling secure Transformer inference in heterogeneous TEE and accelerator systems with dual protection for both model and data. TwinShield offloads ~87% of computation to GPUs and delivers 4.0x - 6.1x speedups over previous approaches across various Transformer models.