Stiefel Manifolds及其对MOORE网络的约束

为了解决多任务问题，RL中的MoE网络架构MOORE通过正交化表示的方式，提升神经网络的表达能力。这种方式相当于对神经网络施加了约束Stiefel Manifold约束。那么，这种约束对神经网络有什么影响呢？

Stiefel Manifold

Manifold是局部为欧几里得空间的拓扑空间。例如：一维流形有线、圆圈，二维流形有平面、球面。其中，欧几里得空间是基本的几何空间，可以直观的理解为物理空间的数学建模。那么，什么是拓扑空间呢？

拓扑空间是一个定义了紧密度的几何空间，但不依赖于具体数值距离而定义的空间。确切的说，拓扑空间是由点的集合构成，这些点的邻域满足特定公理，从而刻画紧密度。拓扑空间在满足极限、连续、连通定义的最通用数学空间，其包含了欧几里得空间和流形。

Stiefel Manifold是欧几里得空间$\mathbb{R}^d$中矩阵$V_s\in\mathbb{R}^{d\times k}$满足$V_s^TV=I_k$，即所有正交的$k$维向量的集合。数学表达为$\mathcal{V}_k(\mathbb{R}^{d})=\{V_s\in\mathbb{R}^{d\times k}:V_s^{T}V_s=I_k\}$

施密特正交化

施密特正交化把线性独立的向量集合$\mathcal{U}=\{u_1,\ldots,u_k:u_i\in\mathbb{R}^d,\forall i\le k\}$变为彼此正交的向量集合$\mathcal{V}=\{v_1,\ldots,v_k:v_i\in\mathbb{R}^d,\forall i\le k\}$，即

$$
\begin{aligned}
v_k=u_k-\sum_{i=1}^{k-1}\frac{\langle v_i,u_k\rangle}{\langle v_i,v_i\rangle}v_i
\end{aligned}\tag{1}
$$

若式(1)变为

$$
\begin{aligned}
v_k+\sum_{i=1}^{k-1}\frac{\langle v_i,u_k\rangle}{\langle v_i,v_i\rangle}v_i=u_k
\end{aligned}\tag{2}
$$

可以理解为$u_k=\alpha_1\cdot v_1+\cdots+\alpha_{k-1}v_{k-1}+v_k$，$\alpha_1\ldots,\alpha_{k-1}$为线性组合系数。对于列满秩的矩阵$\mathcal{U}$，施密特正交化就是进行矩阵的$QR$分解。

由此，$\mathcal{U}=QR$，$Q$为列向量彼此正交的矩阵，$R$为对角元为$1$的上三角矩阵。对于列满秩的矩阵，$QR$分解中$R$为可逆的上三角矩阵，那么$Q=\mathcal{U}R^{-1}$

Stiefel Manifold对MOORE网络的约束

在MOORE中，为了提升策略网络抽取表示的多样性，策略被定义为$\pi(a\vert s,c)=f_{\theta}(\phi(s)\cdot w_c)$，$w_c$为把$k$个正交表示结合为特定任务表示的权重，$f_{\theta}$为基于任务特定表示而生成动作的函数(参数$\theta$为可学习参数)。

确切的说，作者们利用MoE近似$\phi$，从而生成状态$s$的$k$个表示$U_s\in\mathbb{R}^{d\times k}$。RL目标被定义为

$$
\begin{aligned}
\underset{\Theta=\{\phi,\theta\}}{max} J(\Theta) \\
s.t. ~~ \mathbf{h}_{\phi}^{T}(s)\mathbf{h}_{\phi}(s)=\mathbf{I}_k~\forall s\in\mathcal{S}
\end{aligned}\tag{3}
$$

式(3)中$\mathbf{h}$表示专家的输出，即状态的表示；$J(\Theta)$为RL的目标函数，对于策略梯度方法为累积奖励；$\mathbf{I}_k\in\mathbb{R}^{k\times k}$为单位矩阵

然而，在实现上，作者们通过对表示做施密特-正交化的方式，提升表示的多样性。由于目标函数(1)表达的是要求专家的输出为正交表示，这已经偏离了目标函数。

为了方便分析，把优化目标(3)转化为无约束的优化问题

$$
\begin{aligned}
\underset{\Theta=\{\phi,\theta\}}{min} -J(\Theta)+\lambda I_0(\mathbf{h}_{\phi}^{T}(s)\mathbf{h}_{\phi}(s)-\mathbf{I}_k)
\end{aligned}\tag{4}
$$

式(4)中$I_0$为指示函数，只有$u=0$时$I(u)=0$，否则$I(u)=\infty$

那么，正交化约束相当于正则化，有可能提升模型的泛化性，也有可能约束了神经网络的函数空间，限制表达能力。对于MoE网络来说，若每个专家都能很好的学习，那么学习出的表示就应该是多样的，正交化约束可能反而降低性能，这与笔者的实践经验是一致的。

对于显式正交化，笔者认为只是对网络的输出进行了分解，在正反向传播时进行了线性变换，本身不会提升学习表示的多样性。