如何构造支持梯度下降训练的神经网络d维旋转约束层

首先回答你提到的「是否可以直接使用不带非线性激活的普通线性层」的问题：不可以。普通线性层的参数矩阵没有正交约束，会包含缩放、剪切、反射等不属于旋转变换的分量，无法满足你对R严格为旋转矩阵的要求。

以下是三种可落地的实现思路，你可以根据自己的需求选择：

• 方案1：SVD动态正交化（实现最简单，推荐优先尝试） #

  你只需要定义一个形状为(d, d)的普通可学习参数矩阵M，每次前向传播时按以下步骤得到严格的旋转矩阵R：
  1. 对M做奇异值分解，示例为PyTorch写法，其他深度学习框架同理：U, S, Vh = torch.linalg.svd(M)
  2. 得到初始正交矩阵Q = U @ Vh，此时Q是严格正交矩阵，但行列式可能为-1（对应反射变换，不属于旋转）
  3. 行列式修正：如果torch.linalg.det(Q) < 0，将Q的最后一列乘以-1，得到的就是行列式为1的严格旋转矩阵R
     该方案支持自动梯度回传，d>10的场景下计算开销极低，核心代码量仅10行以内。

• 方案2：吉文斯旋转参数化（参数效率最高，严格约束） #

  d维旋转矩阵可以拆解为d*(d-1)/2个二维吉文斯旋转矩阵的乘积，每个二维旋转对应一个可学习的角度参数theta，和你二维场景下的实现逻辑完全一致。以d=10为例，仅需要学习45个参数，远少于方案1的100个参数，且天生保证输出为严格旋转矩阵，不需要额外修正步骤。
  你可以直接基于深度学习框架的张量操作实现旋转矩阵的堆叠计算，也可以用开源社区已经封装好的吉文斯旋转层代码。

• 方案3：正交正则化软约束（运行效率最高，约束宽松） #

  如果你对R的旋转约束不需要100%严格，可以直接用不带激活的普通线性层，额外在损失函数中加入两个正则项即可：

  orth_loss = torch.norm(M.T @ M - torch.eye(d), p='fro')
  det_loss = (torch.linalg.det(M) - 1) ** 2
  total_loss = 原任务损失 + 0.01 * orth_loss + 0.01 * det_loss
  

  该方案不需要在前向传播中做额外矩阵运算，运行效率最高，但无法保证R是严格的旋转矩阵，仅能做到近似约束。

小提示：你用到的R.transpose() @ f(R @ z)结构是旋转等变网络的典型结构，上述三种方案得到的R都可以直接代入该结构使用，梯度可以正常回传训练。

内容的提问来源于stack exchange，提问作者Sean K