京东发布基于张量网络加速的大规模、分布式量子机器学习平台TeD-Q

TeD-Q (Tensor-network enhanced Distributed Quantum) 是京东探索研究院发布的一款基于 Python 可微分编程的量子机器学习软件平台。该平台同时支持求解量子变分算法及量子计算模拟。TeD-Q 为量子计算机与量子模拟器提供统一的架构，在应用层面做到后端无感，用户可以像训练神经网络一样训练量子神经网络。这为量子机器学习以及量子算法的研究提供了一个一站式的解决方案，也为广大量子计算研究人员提供了一个方便易用、性能优异的选择。

TeD-Q 支持量子模拟器与量子硬件两种后端模式，提供可兼容于业界流行的机器学习框架 PyTorch 与 JAX 的接口，实现量子计算与经典深度学习的无缝衔接。其支持即时编译、GPU 平行化加速、反向传播快速自动微分、适用于量子硬件的参数偏移自动微分。

TeD-Q 内建的量子模拟器支持全振幅波函数向量传导模式与快速张量网络缩并模式，且 TeD-Q 可计算出两种模式的计算复杂度，用户可根据量子线路的特点选择适用的模式。TeD-Q 拥有完全独立的自主知识产权，并不需要过度依赖调用第三方核心软件包，内建 Tensor Network 模组，并且可搭配京东探索研究院的发布的张量网络缩并路径搜索工具 JDtensorPath 使用，使张量网络缩并模式的性能达到最优化。

TeD-Q 的架构设计遵从高效易用原则，核心的框架是 QInterpreter 与 Backends。前者负责把用户输入的量子线路描述函数转换成对应的可复用量子线路，后者则是负责高效能的实际运算。与此同时，我们还提供了一些常用的量子线路模板与可视化进度工具还有內建的梯度下降优化器。

图 1. TeD-Q 软件架构

TeD-Q 使用简洁优雅的 API 设计，如图 2 所示，用户可以像编写标准 Python 函数那样创建量子函数。无论使用哪种计算后端，所创建的可执行量子函数在外部看来就如同一般 Python 函数那样，可与现有的经典深度学习框架进行无缝衔接。详细的使用请参考（Qubit_rotation.ipynb 例子与 TeD-Q 说明文档）。

https://github.com/JDEA-Quantum-Lab/TeD-Q/blob/main/examples/Qubit_rotation.ipynb

图 2. 使用 TeD-Q 创建量子函数

图 3. 可视化量子线路编辑

通过底层张量网络引擎赋能与缩并路径优化器的使用，在多数量子线路上可大幅缩减模拟的计算复杂度与内存使用量，从而可以实现传统模拟器无法完成的大规模量子线路模拟。TeD-Q 拥有自己內建的 Tensor Network 模组，其可搭配 cotengra 或 JDtensorPath（推荐）缩并路径优化器使用。

JDtensorPath 是京东探索研究院为 TeD-Q 独立研发的一款配套的缩并路径优化软件，其采用高维图切割、随机搜索、slicing 分割、以及高效多进程管理等手段，能继续增强 TeD-Q 量子线路的仿真性能，使其在量子模拟能力上处于业界领先水平。详情可参照

（Many_body_Localization_2D.ipynb 例子）。

https://github.com/JDEA-Quantum-Lab/TeD-Q/blob/main/examples/Many_body_Localization_2D.ipynb

1.高维图分割，每次分割后补充fake节点并赋予高权重，使得fake节点能够更加均匀的分布在后续切割的两侧，能够找出更优的缩并顺序解。

2.Slicing：

考虑有如下张量：T_abcdefg, T_ahijkl, T_bmnop, T_cqrst, T_duvwx, T_hyzAB, T_iCDEF,它们的 size 分别为 7， 6， 5， 5， 5， 5，5。

假设我们的目标是想通过 slicing 使得它的 size 不大于 4，通过枚举法，我们能得知最佳的解为 “bcdhi”，但枚举法的复杂度为O（N！)，无法应用于复杂的网络。目前主流的 slicing 方式是每次只考虑一个 indice，找出局部最优解。然而这种情况下则很大几率会把并不是最优解的 “a” 也包括进去。JDtensorPath 的随机算法会考虑全局状况，其会先找出所有需要分割的张量,然后通过 python 的集合操作与随机过程，有一定的概率可以找出较优的解。

3.速度与内存消耗：

JDtensorPath 內建有简易的二分型张量缩并树，内存消耗少，读写速度快。同时寻找 slicing 时并不像主流的软件那样需要复制整棵树并进行操作。JDtensorPath 的算法能在不需要修改树的结构与数据的情况下平行化算出不同 slicing 结果的 cost，然后找出最优的slicing方式后再对树的结构内容进行最终修改。此算法与主流的算法相比内存消耗可有数量级的降低。

与国内外业界主流量子模拟方案相比，TeD-Q 的张量网络缩并模式在多种常见量子线路上表现出数量级的加速效应，在某些特定的量子线路上，可达到百万级的加速。在几乎所有的量子线路上，TeD-Q 的波函数向量传导模式也表现出一定的加速效果。图 4 展示了 TeD-Q 与 pennylane、qiskit 在 1-D Quantum Neuronal Sensing of Quantum Many-Body States 例子上的速度对比，在 24 GB 的 GPU 上，pennylane 最多可模拟20个量子比特的线路，而 TedD-Q的张量网络缩并模式能在较快的时间内完成50个量子比特的线路。

图 4. TeD-Q 与 pennylane、qiskit 在 1-D Quantum Neuronal Sensing of Quantum Many-Body States 例子上的速度对比

请参考软件文档与例子学习软件的使用和丰富的实践案例。

本文转自京东探索研究院。

END