Mitsuba 3构造矩阵时梯度丢失：面光源绝对尺度优化失效

我正在Mitsuba 3中以参考图像为目标，用均方误差（MSE）作为损失函数优化面光源的绝对尺度。

初始场景与目标图像：

• 初始场景：
• 目标图像：

通过latent变量生成乘法缩放矩阵的方法可正常运行，但缩放存在指数复合效应，收敛难度大（当缩放值足够大时，代表缩放因子的latent变量远大于1，线性学习率调整无法及时将乘数调回1），不过在优化后的光源超过参考光源尺寸前停止迭代可生效，示例代码如下：

scene = mi.load_file("my_scene.xml", integrator='prb')
params = mi.traverse(scene)
params.update()

# Generate a reference image (area light at correct size)
reference = mi.render(scene, params, spp=1024)

# Disturb the scale by a factor of 0.5
def multiplicative_resize(matrix, factor):
    """Applies local scale around the centroid (multiplicative)."""
    centroid = mi.Vector3f(matrix[0, 3], matrix[1, 3], matrix[2, 3])
    to_origin = mi.Transform4f.translate(-centroid)
    scale_tf = mi.Transform4f.scale(factor)
    from_origin = mi.Transform4f.translate(centroid)
    return from_origin @ scale_tf @ to_origin @ mi.Transform4f(matrix)

light_key = "Light.to_world"
original_matrix = params[light_key].matrix
params[light_key] = multiplicative_resize(original_matrix, 0.5)
params.update()

optimizer = mi.ad.Adam(lr=0.01)
optimizer["latent_scale_factor"] = mi.Float(1.001) # couldn't be exactly 1 otherwise there would be no gradient

for it in range(10):
    # Render
    image = mi.render(scene, params, spp=1024)
    loss = dr.mean(dr.sqr(image - reference))

    dr.backward(loss)
    optimizer.step()

    # Clamp
    scale_val = dr.clamp(optimizer["latent_scale_factor"], 0.1, 2.0)
    optimizer["latent_scale_factor"] = scale_val

    # Update the transform multiplicatively
    current_matrix = params[light_key].matrix
    new_matrix = multiplicative_resize(current_matrix, scale_val)
    params[light_key] = mi.Transform4f(new_matrix)
    params.update()

    print(f"[Multiplicative] Iter {it}, loss = {loss[0]:.6f}")

因此我认为直接设置绝对尺度会让优化更简单，于是尝试让优化器学习匹配面光源的绝对尺度与参考图像，实现代码如下：

def set_area_light_scale(matrix, scale_value):
    """
    Sets the (0,0) and (1,1) entries to `scale_value` 
    directly, ignoring previous scale.
    """
    # Suppose for a rectangular area light, the matrix 
    # defaults to something like:
    #   [ 2  0  0  x ]
    #   [ 0  2  0  y ]
    #   [ 0  0  2  z ]
    #   [ 0  0  0  1 ]
    # Forcibly place `scale_value` in positions (0,0) and (1,1).
    new_mat = dr.llvm.ad.Matrix4f(
        scale_value, matrix[0,1], matrix[0,2], matrix[0,3],
        matrix[1,0], scale_value, matrix[1,2], matrix[1,3],
        matrix[2,0], matrix[2,1], matrix[2,2], matrix[2,3],
        matrix[3,0], matrix[3,1], matrix[3,2], matrix[3,3]
    )
    return mi.Transform4f(new_mat)

scene = mi.load_file("my_scene.xml", integrator='prb')
params = mi.traverse(scene)
params.update()

reference = mi.render(scene, params, spp=1024)

# Disturb the scale absolutely (initially 0.5, e.g.)
light_key = "Light.to_world"
original_matrix = params[light_key].matrix
    
params[light_key] = set_area_light_scale(original_matrix, 0.5)
params.update()

optimizer = mi.ad.Adam(lr=0.01)
optimizer["latent_scale_factor"] = mi.Float(0.5)

for it in range(10):
    image = mi.render(scene, params, spp=1024)
    loss = dr.mean(dr.sqr(image - reference))

    dr.backward(loss)
    optimizer.step()

    # clamp
    scale_val = dr.clamp(optimizer["latent_scale_factor"], 0.1, 2.0)
    optimizer["latent_scale_factor"] = scale_val

    # forcibly set the scale
    current_tf = params[light_key].matrix
    params[light_key] = set_area_light_scale(current_tf, scale_val)
    params.update()

    print(f"[Absolute Scale] Iter {it}, loss = {loss[0]:.6f}")

新函数可正常更新面光源尺寸，但问题在于现在每次迭代的损失值完全相同，没有下降。我是否破坏了梯度？问题出在哪里？

TL;DR #

• 乘法缩放变换方法梯度流动正常，但迭代中缩放变化呈指数型；
• 绝对尺度方法每次迭代损失值完全相同；
• 希望找出绝对尺度对应的梯度无法降低损失的原因及修复方法。

核心原因：梯度传递链被切断 #

你的绝对尺度方法中，每次迭代更新时都忽略了当前变换矩阵的梯度信息，直接用current_tf覆盖并强制设置缩放值，导致梯度无法在优化器的latent变量和场景参数之间形成连续传递：

1. 从params[light_key].matrix获取的current_tf是上一次迭代后被强制设置的固定值，没有保留梯度信息；
2. 优化器更新的latent_scale_factor虽然有梯度，但在更新场景参数时，完全丢弃了之前变换矩阵的梯度，相当于每次都在重新初始化变换，梯度无法累积传递。

修复方案：保留梯度关联的绝对尺度优化 #

不需要每次从params读取当前矩阵，而是直接基于原始矩阵生成变换，让latent变量和场景变换始终保持梯度依赖：

修改后的代码 #

def set_area_light_scale(original_matrix, scale_value):
    """
    基于原始矩阵设置绝对尺度，保留梯度关联
    """
    new_mat = dr.llvm.ad.Matrix4f(
        scale_value, original_matrix[0,1], original_matrix[0,2], original_matrix[0,3],
        original_matrix[1,0], scale_value, original_matrix[1,2], original_matrix[1,3],
        original_matrix[2,0], original_matrix[2,1], original_matrix[2,2], original_matrix[2,3],
        original_matrix[3,0], original_matrix[3,1], original_matrix[3,2], original_matrix[3,3]
    )
    return mi.Transform4f(new_mat)

scene = mi.load_file("my_scene.xml", integrator='prb')
params = mi.traverse(scene)
params.update()

reference = mi.render(scene, params, spp=1024)

light_key = "Light.to_world"
original_matrix = params[light_key].matrix  # 仅保存一次原始矩阵

optimizer = mi.ad.Adam(lr=0.01)
optimizer["latent_scale_factor"] = mi.Float(0.5)  # 初始扰动值

for it in range(10):
    # 直接用原始矩阵+当前latent值生成变换，保留梯度链
    params[light_key] = set_area_light_scale(original_matrix, optimizer["latent_scale_factor"])
    params.update()

    image = mi.render(scene, params, spp=1024)
    loss = dr.mean(dr.sqr(image - reference))

    dr.backward(loss)
    optimizer.step()

    # 对latent值做范围限制
    optimizer["latent_scale_factor"] = dr.clamp(optimizer["latent_scale_factor"], 0.1, 2.0)

    print(f"[Absolute Scale] Iter {it}, loss = {loss[0]:.6f}")

关键修改点 #

1. 固定原始矩阵：不再每次从params读取当前矩阵，始终基于最初的original_matrix生成变换，确保latent变量的梯度能直接传递到变换矩阵；
2. 调整更新顺序：先根据当前latent值更新场景参数，再渲染计算损失，保证每次迭代的变换都是基于最新的优化器变量；
3. 移除冗余操作：避免破坏梯度传递链，让优化器的latent变量和场景变换始终保持直接依赖关系。

额外优化建议 #

• 如果面光源非轴对齐，建议先提取原始矩阵中的旋转和平移分量，再单独缩放光源的局部轴，避免破坏光源朝向；
• 绝对尺度优化的梯度通常更稳定，可尝试增大学习率（如lr=0.05）加速收敛。

内容的提问来源于stack exchange，提问作者Anson Savage