使用 TVM 实现高效的隐私保护机器学习
这篇文章介绍了 Myelin，一个用于在可信硬件飞地中进行隐私保护机器学习的框架，以及 TVM 如何使 Myelin 变得快速。 核心思想是 TVM 与其他流行的 ML 框架不同，它将模型编译成轻量级、优化且无依赖的库，这些库可以容纳在资源受限的飞地中。
尝试创建隐私保护的 ML 模型！ 查看 TVM 仓库中提供的示例代码。
动机：隐私保护机器学习
机器学习模型受益于大型且多样化的数据集。 不幸的是，使用此类数据集通常需要信任中心化的数据聚合器或计算提供商。 对于医疗保健和金融等敏感应用，这是不可取的，因为它可能会泄露患者隐私或泄露商业机密。 安全和隐私保护计算的最新进展，包括可信执行环境和差分隐私，为互不信任的各方提供了一种有效的方法来训练机器学习模型，而不会泄露训练数据。 我们使用 TVM 使隐私保护 ML 框架变得快速。
用例
- 私有 MLaaS：云提供商在您的数据上运行其架构。 您获得模型输出，您的数据保持私有，云提供商知道您无法窃取模型。
- 值得信赖的 ML 竞赛：您在竞赛数据上训练模型。 竞赛组织者将私有测试数据发送到您的模型并获得可验证的准确性报告。 您的模型保持安全，直到组织者决定购买它。 其他参与者不能通过在测试数据上进行训练来作弊。
- 在共享私有数据上进行训练：您（研究人员）想在几家医院的数据上训练模型。 直接共享过于复杂。 相反，让“可信第三方”训练隐私保护模型。
- 区块链上的机器学习
背景
可信执行环境
可信执行环境 (TEE) 本质上允许远程用户在另一个人的机器上可证明地运行代码，而不会向硬件提供商泄露计算。
更技术性地说，TEE 提供了一个安全的飞地，其中包含隔离/加密的内存和 CPU 寄存器； 以及可信的随机源。 TEE 还可以发送已加载代码的签名证明，以便远程用户可以验证飞地是否已正确加载。 此过程称为远程证明，可用于建立到飞地的安全通信通道。 然后，远程用户可以使用私钥、模型参数和训练数据等机密来配置它。
与纯加密方法（如 安全多方计算 (MPC) 和 全同态加密 (FHE)）相比，TEE 快几个数量级，并且支持通用计算（即不仅仅是算术运算）。 也许唯一的缺点是对硬件信任根（烧录到处理器中的密钥）和加载的软件的额外信任假设。
尽管存在信任假设，但 TEE 技术正变得越来越普及，并在实际的隐私保护中发挥着重要作用。 实际上，通用 TEE 已经存在于商品硬件中，例如 Intel SGX 和 ARM TrustZone。 此外，完全开源的 Keystone 飞地 也即将问世。
差分隐私
差分隐私 (DP) 提供了正式保证，即在相似数据集上训练的模型是无法区分的。 非正式地，用户的隐私不会因选择向模型贡献数据而受到损害。
换句话说，给定算法在仅在一个记录不同的两个数据集上的输出，差分隐私限制了对手可以确定哪个数据集的概率上限。 可以使用一种机制使算法成为 DP，该机制向算法的输出添加噪声。 噪声量根据输出在多大程度上取决于任何特定输入进行校准。 如果您熟悉假设检验，如果结果 A 和 B 各自具有 0.5 的概率，则应用 DP 机制就像与概率分布卷积一样：隐私在于误报率和误报率。 由于深度学习模型往往具有良好的泛化能力，因此噪声量通常小于预期。
在 TEE 中运行 DP 训练算法可确保 DP 机制得到忠实应用。
使用 TVM 实现高效的隐私保护机器学习
使用 TEE 的主要挑战之一是，在其中运行的代码无法访问不受信任的操作系统。 这意味着受信任的软件无法创建线程或执行 I/O。 实际上，结果是像 OpenBLAS 这样的数值库——更不用说像 PyTorch 和 TensorFlow 这样的框架——无法直接在飞地中运行。
TEE 实际上具有类似于资源受限硬件加速器的编程模型。 这正是 TVM 的用途！ 在隐私工作流程中，用户首先以高级图规范语言定义整个训练图。 然后 TVM 编译模型并输出一个静态库，其中包含可以轻松加载到 TEE 中的优化数值内核。 由于内核是自动生成的并且具有严格的边界检查，因此它们暴露了低攻击面。 它们由轻量级内存安全的 Rust 运行时支持，该运行时也很容易审查其安全性和正确性。
当然，安全在实际应用时最有用。 幸运的是，飞地中的 TVM 模块具有与基于 CPU 的本地训练相当的性能。 通过使用不受信任的运行时协调线程，单个 TVM 飞地可以充分利用其主机资源的资源。 此外，不难想象一个安全参数服务器，它可以协调整个数据中心的启用飞地的机器。
TVM 还为更精细地优化隐私保护算法提供了机会。 事实上，当使用差分隐私时，其细粒度的调度功能可以加速。 例如，最严格的 DP 边界可以从剪裁每个训练示例的梯度并向每个梯度添加噪声 [1] 获得。 在自动求导框架中，这需要为小批量中的每个示例转发模型（尽管只需要一个反向传递）[2]。 然而，使用 TVM，每个示例的梯度裁剪非常简单：不是将每个权重更新调度为批次和特征维度上的单个缩减，而是将缩减分为两个。 特征上的缩减之后是裁剪和添加噪声，然后最终结果在示例上求和以获得权重更新。 因此，TVM 允许应用差分隐私，而不会引入大于该技术所需开销的开销。 此外，如果真的想变得非常花哨，可以融合裁剪和添加噪声操作，并就地应用它们，以进一步减少延迟和内存使用量。
有关实际工作负载的基准测试，请参阅技术报告 多源私有数据上的高效深度学习。 当然，也可以随意在 TVM SGX 示例中试用该框架。
结论
下一代学习系统将由隐私开启。 随着 TEE 技术变得更容易理解和更广泛地可用，将其用作隐私保护机器学习和分析的资源是有意义的。 TVM 已做好充分准备，以促进研究和部署中此用例的开发。
简介和致谢
Nick 是加州大学伯克利分校 Dawn Song 教授实验室的博士生。 他的研究兴趣在于共享私有数据上的 ML 的一般领域，但这实际上只是一个借口，可以摆弄 Rust、安全监控器、硬件飞地和像 TVM 这样的编译器。
感谢 Tianqi Chen 的代码审查！
参考文献
[1] 具有差分隐私的深度学习
[2] 高效的每个示例梯度计算