TT分解与网络张量化

TT分解

思路

将一个张量表示为一种运算，即多个低维张量的乘积。只需储存分解后低维张量，进而压缩占用的储存空间。

数学表示与推导

对于一个 $d$ 维张量
$\mathcal{A} \in \mathbb{R}^{p_{1} \times p_{2} \times \ldots \times p_{d}}$ ， $p_i$ 表示第 $i$ 维的长度。可以分解为多个低维张量的乘。

$\widehat{\mathcal{A}}\left(l_{1}, l_{2}, \ldots, l_{d}\right) \stackrel{T T F}{=} \mathcal{G}_{1}\left(l_{1}\right) \mathcal{G}_{2}\left(l_{2}\right) \ldots \mathcal{G}_{d}\left(l_{d}\right)$
规定，
$\mathcal{G}_{k} \in \mathbb{R}^{p_{k} \times r_{k-1} \times r_{k}}, l_{k} \in\left[1, p_{k}\right] \forall k \in[1, d] ,r_0=1, r_d=1$

image

用TT分解压缩全连接层

对于一个全连接层，数学表示为： $\boldsymbol{y}=\boldsymbol{W} \boldsymbol{x}+\boldsymbol{b}$ 。要压缩全连接层，重点在于压缩矩阵 $\boldsymbol{W}$ 和偏差向量 $\boldsymbol{b}$ 。

矩阵和向量的TT分解形式

向量的TT分解

向量的TT分解实际上是将向量按张量的方式存储，并未压缩

对于一个 $N$ 长的向量，将 $N$ 表示为 $N=n_1\cdot n_2 \cdots n_d$ ，则将向量分解为一个 $d$ 维的张量，第 $i$ 维的长度就是对应的 $n_i$ 。

使用一个向量 $\boldsymbol{\mu}(\ell)=(\mu_1(\ell),\cdots,\mu_d(\ell))$ 来联系张量和向量中的元素，有 $\mathcal{B}(\boldsymbol{\mu}(\ell))=b_{\ell}$ ，其中 $\mu_k(\ell)$

矩阵的TT分解

将矩阵映射到张量，再做TT分解

与向量类似，对于 $M\times N$ 的矩阵，分别将 $M,N$ 表示为 $d$ 个数的相乘。使用向量 $\boldsymbol{\nu}(t)=\left(\nu_{1}(t), \cdots, \nu_{d}(t)\right)$ 和 $\boldsymbol{\mu}(\ell)=\left(\mu_{1}(\ell), \cdots, \mu_{d}(\ell)\right)$ 来联系张量和矩阵中 $(t,\ell)$ 位置的元素。张量中第 $k$ 维的大小为 $m_kn_k$ 。

于是矩阵 $W$ 中的某个元素可以按下面的规则对应到 $d$ 维张量中的元素，再对张量TT分解，得到矩阵的TT形式。
$W(t, \ell)=\mathcal{W}\left(\left(\nu_{1}(t), \mu_{1}(\ell)\right), \ldots,\left(\nu_{d}(t), \mu_{d}(\ell)\right)\right)=\boldsymbol{G}_{1}\left[\nu_{1}(t), \mu_{1}(\ell)\right] \ldots \boldsymbol{G}_{d}\left[\nu_{d}(t), \mu_{d}(\ell)\right]$

TT-Layer

完成了向量和矩阵的TT分解后，为便于将全连接层用TT分解形式表示，将输入向量 $\boldsymbol{x}$ 和输出向量 $\boldsymbol{y}$ 同样用TT分解表示为两个 $d$ 维张量。最终，一个全连接层可以表示为：
$\mathcal{Y}\left(i_{1}, \ldots, i_{d}\right)=\sum_{j_{1}, \ldots, j_{d}} \boldsymbol{G}_{1}\left[i_{1}, j_{1}\right] \ldots \boldsymbol{G}_{d}\left[i_{d}, j_{d}\right] \mathcal{X}\left(j_{1}, \ldots, j_{d}\right)+\mathcal{B}\left(i_{1}, \ldots, i_{d}\right)$

TT-Layer的随机梯度下降

直接计算损失函数 $L$ 关于 $W$ 的梯度要内存占用大，使用TT分解形式的梯度计算。

将序列分为 $\boldsymbol{i}^-_k:=\left(i_{1}, \ldots, i_{k-1}\right), \boldsymbol{i}_{k}^{+}:=\left(i_{k+1}, \ldots, i_{d}\right), \boldsymbol{i}=\left(\boldsymbol{i}_{k}^{-}, i_{k}, \boldsymbol{i}_{k}^{+}\right)$ 三部分，那么 $\mathcal{W}$ 也被分为了三部分：
$\boldsymbol{P}_{k}^{-}\left[\boldsymbol{i}_{k}^{-}, \boldsymbol{j}_{k}^{-}\right]:=\boldsymbol{G}_{1}\left[i_{1}, j_{1}\right] \ldots \boldsymbol{G}_{k-1}\left[i_{k-1}, j_{k-1}\right]$
$\boldsymbol{P}_{k}^{+}\left[\boldsymbol{i}_{k}^{+}, \boldsymbol{j}_{k}^{+}\right]:=\boldsymbol{G}_{k+1}\left[i_{k+1}, j_{k+1}\right] \ldots \boldsymbol{G}_{d}\left[i_{d}, j_{d}\right]$
则全连接层又可写为：
$\mathcal{Y}(i)=\mathcal{Y}\left(i_{k}^{-}, i_{k}, i_{k}^{+}\right)=\sum_{j_{k}^{-}, j_{k}, j_{k}^{+}} P_{k}^{-}\left[i_{k}^{-}, j_{k}^{-}\right] G_{k}\left[i_{k}, j_{k}\right] P_{k}^{+}\left[i_{k}^{+}, j_{k}^{+}\right] \mathcal{X}\left(j_{k}^{-}, j_{k}, j_{k}^{+}\right)+\mathcal{B}(i)$
计算TT分解形式下 $L$ 关于 $\mathcal{W}$ 的梯度 $\frac{\partial \mathcal{Y}(\boldsymbol{i})}{\partial \boldsymbol{G}_{k}\left[\tilde{i}_{k}, \tilde{j}_{k}\right]}$ 时只需关注 $i_{k}=\tilde{i}_{k}$ 和 $j_{k}=\tilde{j}_{k}$ 而其他位置都为0， $(\tilde{i}_{k},\tilde{j}_{k})$ 是 $\mathcal{W}$ TT分解后k维张量上的坐标。

$\mathcal{Y}$ 又是 $G_k$ 的线性函数，一求导， $G_k$ 项和常数项 $\mathcal{B}$ 就没了，只剩下:
$\sum_{j_{k}^{-}, j_{k}^{+}} (P_{k}^{-}\left[i_{k}^{-}, j_{k}^{-}\right]) ^T( P_{k}^{+}\left[i_{k}^{+}, j_{k}^{+}\right] )^T\mathcal{X}\left(j_{k}^{-}, \tilde{j_{k}}, j_{k}^{+}\right)$