题目重点：矩阵乘法可结合但不可交换，梯度也能用仿射变换做并行scan。
题意与假设
这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。
- A为n个d×d实数方阵，P_i=A_0…A_i；训练L=Σ_i〈G_i,P_i〉，G为与每个前缀对应的上游梯度。只最后乘积有loss是其特例。
- 原地forward会覆盖输入；反向需要原始A，除非调用方保留或能重算。常数额外空间按n计算：允许O(d²)工作矩阵，排除只读A/G和必须返回的n个梯度矩阵，不排除所有矩阵存储。
- 代码NumPy循环模拟scan阶段，并非实际GPU并行kernel。数组必须浮点或兼容dtype；原地数组dtype会影响赋值精度。
解题路线
- 顺序forward保存累乘，P_i=P_{i−1}A_i；原地直接写A[i]=A[i−1]@A[i]，out-of-place复制后同样处理。遇n0输出空序列，不能索引A[0]。
- 反向令D_i=∂L/∂P_i，从后往前D_i=G_i+D_{i+1}A_{i+1}ᵀ，∂L/∂A_i=P_{i−1}ᵀD_i，P_{−1}=I。每步将back右乘A_iᵀ为下一步准备。无需逆矩阵，A奇异时也成立。
- 常数额外空间版每次从A0重算P_{i−1}，以O(n²d³)计算换取O(d²)工作空间。存全部prefix可将反向工作降为O(nd³)，代价O(nd²)。不能同时给重算实现贴线性时间标签。
- Hillis–Steele forward每轮距离step倍增，用旧数组previous[i−step]@previous[i]，顺序不能反。每轮只读旧阶段，以免同一轮传播多次。反向把f_i(B)=B M_i+C_i，其中M_i=A_{i+1}ᵀ、C_i=G_i；组合f_i∘f_j得到M=M_jM_i、C=C_jM_i+C_i，可做后缀scan，再用prefixᵀ计算梯度。
手动走一遍
取A0=[[1,1],[0,1]]，A1=[[2,0],[0,3]]，A0A1=[[2,3],[0,3]]，A1A0=[[2,2],[0,3]]；顺序颠倒结果不同。仅最终loss时把所有G设0，仅G_{n−1}给定；全前缀loss则每步都要加G_i，不能只从末尾传一次。1000组随机/奇异矩阵用PyTorch autograd核对顺序forward、scan forward、重算backward及仿射scan backward。
Python 参考实现
下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。
代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。
import numpy as np
def prefix_products(matrices,in_place=False):
a=matrices if in_place else np.array(matrices,copy=True)
if a.ndim!=3 or a.shape[1]!=a.shape[2]:raise ValueError('square matrix sequence')
for i in range(1,len(a)):a[i]=a[i-1]@a[i]
return a
def prefix_scan(matrices):
# Hillis-Steele: every stage reads only the previous stage.
out=np.array(matrices,copy=True);step=1
if out.ndim!=3 or out.shape[1]!=out.shape[2]:raise ValueError('square matrix sequence')
while step<len(out):
previous=out.copy()
for i in range(step,len(out)):out[i]=previous[i-step]@previous[i]
step*=2
return out
def all_prefix_gradients(matrices,upstream):
# Loss L=sum_i <upstream[i], A0...Ai>. Constant auxiliary matrices,
# excluding read-only inputs and the required gradient output. No inverses.
a=np.asarray(matrices);g=np.asarray(upstream)
if a.ndim!=3 or a.shape[1]!=a.shape[2] or g.shape!=a.shape:raise ValueError('shape')
n,d,_=a.shape;result=np.zeros_like(a,dtype=np.result_type(a,g,float));back=np.zeros((d,d),dtype=result.dtype)
for i in range(n-1,-1,-1):
back=back+g[i]
left=np.eye(d,dtype=result.dtype)
for j in range(i):left=left@a[j]
result[i]=left.T@back
back=back@a[i].T
return result
def final_product_gradients(matrices,upstream):
a=np.asarray(matrices)
if a.ndim!=3 or a.shape[1]!=a.shape[2]:raise ValueError('shape')
g=np.zeros_like(a,dtype=np.result_type(a,upstream,float))
if len(a):g[-1]=upstream
return all_prefix_gradients(a,g)
def scan_prefix_gradients(matrices,upstream):
a=np.asarray(matrices);g=np.asarray(upstream)
if a.ndim!=3 or a.shape[1]!=a.shape[2] or g.shape!=a.shape:raise ValueError('shape')
n,d,_=a.shape
if not n:return np.zeros_like(a,dtype=np.result_type(a,g,float))
prefix=prefix_scan(a);identity=np.eye(d,dtype=np.result_type(a,g,float))
# f_i(B)=B@M_i+C_i, D_i=f_i(D_{i+1}), terminal D_n=0.
transforms=np.stack([a[i+1].T if i+1<n else identity for i in range(n)])
constants=np.array(g,copy=True,dtype=identity.dtype);step=1
while step<n:
old_m=transforms.copy();old_c=constants.copy()
for i in range(n-step):
transforms[i]=old_m[i+step]@old_m[i]
constants[i]=old_c[i+step]@old_m[i]+old_c[i]
step*=2
result=np.empty_like(constants)
for i in range(n):result[i]=(prefix[i-1].T if i else identity)@constants[i]
return result
复杂度
经典方阵乘O(d³)。顺序forward O(nd³)；常数额外矩阵backward O(n²d³)，工作O(d²)且输出O(nd²)。Hillis–Steele O(nlogn·d³)总工作、O(logn)阶段、O(nd²)内存；实际阶段跨度还含矩阵kernel成本。
易错点与追问
容易失分的边界
- 将backward写成乘inverse(A)会在奇异矩阵失败。
- scan更新读当前out而非上一阶段会破坏分块组合。
- 原地输入已变prefix时再当原始A求梯度，会得到错误结果。
继续深挖
- Blelloch scan怎样把总工作降到O(n)，浮点结合误差有哪些变化？
- 非方阵或矩阵链尺寸不同还能否统一做scan，怎样表达形状约束？
