Efficient Finite Initialization with Partial Norms for Tensorized Neural Networks and Tensor Networks Algorithms

本文介绍了两种高效的算法，用于初始化张量化神经网络及通用张量网络算法，这些算法通过迭代利用子网络的偏弗罗贝尼乌斯范数和正线性逐元素和，在实现有限归一化的同时复用中间计算结果。

要点

想象一下，你正试图用成千上万个微小的乐高积木搭建一座巨大而复杂的塔。这座塔代表“张量网络”，一种用于执行预测天气或理解人类语言等复杂任务的特殊计算机大脑。

本文描述的问题在于，当你试图开始搭建这座塔时会发生什么。如果你只是随手抓起一把积木开始随机堆叠，可能会发生两件坏事：

1. 爆炸：塔生长得如此迅速，以至于变得无限高，导致计算机崩溃，因为数值变得过大而无法容纳。
2. 消失：塔收缩得如此迅速，以至于变得不可见，变成一个计算机甚至无法察觉的微小斑点。

本文介绍了两种巧妙的“智能启动”方法，以确保无论你有多少块积木（或多少层），塔都能以完美的尺寸开始构建。

两种智能启动方法

作者根据你所使用的“积木”类型，制定了两种不同的方案。

1. “弗罗贝尼乌斯”方法（适用于通用积木）

这相当于检查你正在生长的塔的总重量。

• 工作原理：你不是先搭建整座塔然后发现它太重，而是分小块搭建。在添加几层之后，你会暂停并称量该特定部分的重量。
• 修正措施：如果该部分变得过重（过大），你就轻轻地将该部分中的每一块积木缩小一点点。如果太轻，就将其稍微放大。
• 神奇之处：本文的秘诀在于，每次犯错时你不必从头开始。如果你修正了前三层，当你继续构建第四层时，前三层将保持修正后的状态。你重用了之前的工作，从而节省了时间和能量。

2. “线性”方法（仅适用于正数积木）

这种方法适用于每一块积木上都标有正数的塔（就像数苹果，你不能有负数个苹果）。

• 工作原理：你不是称量塔的重量，而是简单地计算当前部分中苹果的总数。
• 修正措施：如果你拥有的苹果太多，就将其缩小；如果太少，就将其放大。
• 特殊之处：本文发现，这种“计数”方法往往比“称重”方法更平滑、更高效，特别是对于非常大的塔而言。它沿着一条笔直、可预测的线生长，而不是狂野的曲线。

为什么这很重要（根据本文）

作者在测试了不同形状的塔（称为张量链和 PEPS）后发现：

• 扩展性良好：无论你拥有只有 5 层的小塔，还是有 30 层的巨型塔，这些方法都能防止数值爆炸或消失。
• 高效：通过重用前一步骤的计算结果，计算机无需重复进行数学运算。
• 实用：他们甚至制作了一个免费的开源工具（一个 Python 函数），以便任何人都可以使用这些“智能启动”方案来构建自己的 AI 模型，而无需担心数值失控。

本文未声称的内容

重要的是要遵循作者实际所说的话：

• 他们并未声称这能使 AI 在长期内变得更聪明或更准确；他们仅修正了起点。
• 他们并未在诊断疾病或驾驶汽车等具体的现实世界问题上测试此方法。他们在网络结构本身的数学层面上进行了测试。
• 他们并未声称这适用于每一种可能的 AI 模型类型，仅适用于使用这些特定“张量网络”结构构建的模型。

简而言之，本文提供了一种可靠的方法，用于在开始播放音乐之前调节巨型扬声器系统的音量旋钮，确保声音既不会大到无法听见，也不会小到无法察觉，同时还能避免你每次转动旋钮时都要重置系统。

技术摘要

以下是论文《用于张量化神经网络与张量网络算法的高效有限范数初始化》的详细技术总结。

1. 问题陈述

张量化神经网络（TNNs）和通用张量网络（TN）算法（例如矩阵乘积态/TT、投影纠缠对态/PEPS）面临一个关键的初始化挑战，即张量值的爆炸或消失。

• 机制：在一个包含 $N$ 个节点的张量网络中，最终表示的张量元素是 $N$ 个核心元素的乘积。如果采用标准分布（如高斯分布）进行初始化，最终元素的幅度会随着节点数量（$N$）和键维（$b$）呈指数级增长。
  • 爆炸：数值变得过大，超出浮点数表示范围（变为无穷大）。
  • 消失：数值变得过小（下溢为零）。
• 现有方案的局限性：
  • 完全收缩：由于内存呈指数级增长，对于大型层而言，计算完整张量以进行重缩放是不可能的。
  • 启发式重缩放：简单地更改初始化超参数（均值/标准差）通常效率低下，且需要反复试验。
  • 酉矩阵/恒等矩阵方法：现有方法（如 Haar 测度、恒等矩阵加噪声）通常仅适用于特定架构（如 MPS），无法很好地推广到 PEPS 或张量流矩阵（TT-M）等复杂结构。

2. 方法论

作者提出了两种迭代算法，利用部分范数计算来对网络进行归一化，而无需计算完整张量。核心创新在于在迭代过程中重用中间计算结果。

A. 弗罗贝尼乌斯张量网络重正化 (FTNR)

• 目标：具有实数条目的通用张量网络。
• 度量：使用弗罗贝尼乌斯范数 ($||A||_F = \sqrt{\sum |a_{ij}|^2}$)。
• 机制：
  1. 部分平方范数：算法不收缩整个网络，而是计算由前 $n$ 个节点组成的子网络的平方弗罗贝尼乌斯范数 ($||A_n||_F^2$)。
  2. 迭代修正：检查部分范数是否在目标容差范围内。
     • 如果部分范数为 $\infty$（发散）或 $0$（消失），算法会对涉及该子网络的节点应用缩放因子。
     • 如果范数为有限值但超出目标范围，则应用特定的缩放因子 $r = (S_n / S^*_n)^{1/(2n)}$。
  3. 效率：关键在于，在归一化步骤之后，中间收缩张量被保存。在下一次迭代中，算法从最后一个成功归一化的节点开始，而不是从节点 1 重新开始，从而显著降低了计算成本。
  4. 处理发散：如果某一步导致 $\infty$ 或 $0$，则应用一个随机缩放因子（数量级）以打破循环并重试。

B. 线性张量网络重正化 (LTNR)

• 目标：表示条目为非负的张量网络（例如概率分布、特定量子态）。
• 度量：使用正线性逐元素和 ($||A||_L = \sum a_{ij}$)。
• 机制：
  • 与 FTNR 类似，但使用元素之和而非平方和。
  • 计算成本低于弗罗贝尼乌斯范数，因为它涉及与全 1 向量 ($\mathbf{1}$) 的收缩，而不是共轭副本。
  • 缩放因子：$r = (L_n / L^*_n)^{1/n}$。
  • 该方法特别有效，因为线性和随条目数量线性增长，而弗罗贝尼乌斯范数随平方和的平方根增长，通常能带来更平滑的收敛。

3. 主要贡献

1. 新颖的初始化协议：引入了 FTNR 和 LTNR，使得可以在不发生内存溢出的情况下初始化任意大规模的张量网络。
2. 部分范数策略：利用部分范数（子网络）允许在完整张量形成之前进行归一化检查，从而在“爆炸”发生前予以预防。
3. 中间计算重用：算法存储临时的收缩张量，使得归一化过程可以从失败点恢复，而不是从头开始，从而优化了计算效率。
4. 通用性：这些方法适用于各种架构，包括张量流（TT）、张量流矩阵（TT-M）和 PEPS，涵盖了通用和非负条目场景。
5. 开源实现：作者提供了 Python/PyTorch 实现和 Streamlit 演示，使该方法易于实际使用。

4. 实验结果

作者在具有不同节点数（$N$）、物理维度（$p$）和键维（$b$）的 TT 和 TT-M 层上测试了这些算法。

• 随节点数（$N$）的扩展：
  • 对于小型网络（$N < 10$），不需要归一化步骤。
  • 对于中等规模（$N \approx 27$），通常只需要一步。
  • 对于非常大的 $N$，所需步骤数呈指数级增加，但算法在标准初始化会失败的情况下成功收敛。
• 随物理维度（$p$）的扩展：
  • 对于大 $p$，所需步骤数也呈现类似的指数增长，但 LTNR 算法通常比 FTNR 需要更少的步骤。
• 随键维（$b$）的扩展：
  • 未观察到步骤数对 $b$ 有实质性依赖，这可能是因为算法根据计算出的部分范数进行自适应缩放。
• 比较：**LTNR（线性）**方法始终优于 FTNR，需要更少的迭代次数。这归因于正线性和相比弗罗贝尼乌斯范数的二次性质具有更平滑的缩放行为。

5. 意义与未来应用

• 实现大规模 TNNs：这项工作消除了训练张量化深度学习模型的主要瓶颈，使得使用以前因数值不稳定而无法训练的包含数百个节点的层成为可能。
• 超越深度学习：这些方法适用于任何需要具有相似量级非零元素的张量收缩的算法，例如：
  • 量子机器学习：将经典模型压缩为受量子启发的架构。
  • 物理模拟：使用张量化物理信息神经网络求解微分方程（如热方程、流体动力学）。
  • 组合优化：确定优化问题中的超参数和衰减因子。
• 未来方向：作者建议未来的研究集中在减少所需步骤的数量、分析不同层类型的复杂度扩展，以及将这些方法应用于量子机器学习层。

总之，本文提供了一种稳健、高效且通用的解决方案，解决了张量网络中的初始化问题，促进了复杂、高维模型在经典和受量子启发的机器学习中的部署。