[论文解读] tntorch: Tensor Network Learning with PyTorch

Mikhail Usvyatsov, Rafael Ballester‐Ripoll+1 more|arXiv (Cornell University)|Jun 22, 2022

Tensor decomposition and applications被引用 10

一句话总结

tntorch 是一个与 PyTorch 集成的张量网络库，通过统一的、类似 PyTorch 的 API，将多种低秩张量分解（CP、Tucker 和 Tensor Train）整合在一起。它支持自动微分、GPU 加速以及无缝的可微操作，使用户能够高效地学习、操作和压缩高维张量，同时支持混合格式网络和高级功能，如交叉近似和可微张量代数。

ABSTRACT

We present tntorch, a tensor learning framework that supports multiple decompositions (including Candecomp/Parafac, Tucker, and Tensor Train) under a unified interface. With our library, the user can learn and handle low-rank tensors with automatic differentiation, seamless GPU support, and the convenience of PyTorch's API. Besides decomposition algorithms, tntorch implements differentiable tensor algebra, rank truncation, cross-approximation, batch processing, comprehensive tensor arithmetics, and more.

研究动机与目标

• 将多种张量分解格式（CP、Tucker、TT）统一到一个用户友好的接口中，抽象底层表示差异。
• 提供与 PyTorch 自动微分和 GPU 加速完全兼容的支持，实现端到端的可微张量学习。
• 支持在单一计算图中组合 CP、Tucker 和 TT 组件的混合张量网络架构。
• 实现高效的可微操作，如交叉近似、秩截断和张量代数，用于机器学习和数据分析。
• 提供一个灵活且可扩展的框架，支持高级用例，如张量补全、模型压缩以及通过黑箱采样实现的代理建模。

提出的方法

• 将所有张量分解封装到一个统一的 `Tensor` 类中，采用与 PyTorch 兼容的 API，实现与深度学习工作流的无缝集成。
• 将张量表示为低秩张量网络，其中每个节点为一个核心或因子，通过张量收缩执行操作，并自动计算梯度。
• 通过允许 CP、Tucker 和 TT 组件在单个张量对象中以任意组合共存并组合，支持混合格式张量网络。
• 使用 maxvol 算法实现可微的交叉近似，从黑箱张量函数中采样关键索引，实现在无需完整存储张量情况下的内存高效学习。
• 提供专用类如 `TTMatrix` 和 `CPMatrix`，以高效实现 TT 和 CP 格式下的矩阵-向量乘法，避免完全解压张量。
• 利用 PyTorch 的 autograd 系统，支持基于梯度的优化，包括全局优化和黎曼优化，用于学习低秩张量结构。

实验结果

研究问题

• RQ1 如何将多种张量分解格式（CP、Tucker、TT）统一到一个可组合的接口中，同时保持可微性和 GPU 支持？
• RQ2 交叉近似能否实现可微化，并集成到深度学习流水线中，以实现对大型或内存中张量的高效、基于梯度的学习？
• RQ3 混合张量网络架构（如 CP-TT-Tucker 混合）在机器学习任务中能在多大程度上提升表达能力和效率？
• RQ4 如何高效地实现可微且可扩展的张量代数操作，如卷积、填充和秩截断？
• RQ5 一个原生支持 PyTorch 的库能否在性能和灵活性上与专用张量库相媲美，同时保持易用性和可扩展性？

主要发现

• tntorch 成功地将 CP、Tucker 和 TT 分解统一到一个 `Tensor` 类中，全面支持自动微分和 GPU 计算。
• 通过 maxvol 算法实现了可微的交叉近似，使用户能够从黑箱函数中高效学习 TT 张量，而无需存储完整张量。
• 原生支持混合张量网络架构（如 CP-TT-Tucker），可使用标准 PyTorch 操作端到端构建和优化。
• 可微张量代数操作（包括逐元素函数、卷积和秩截断）可高效计算，内存开销极低。
• `TTMatrix` 和 `CPMatrix` 的实现可在不解压张量的情况下实现快速矩阵-向量乘法，时间复杂度与核心数量呈线性关系。
• tntorch 支持高级功能，如敏感性分析、统计矩计算以及从 TT 压缩分布中采样，推动其在概率建模和优化中的广泛应用。