Iskra: A System for Inverse Geometry Processing

一句话总结

iskra 是一个建立在 PyTorch 之上的命令式、张量化系统，用伴随法（adjoint method）自动地对几何处理求解器的”解”求导，让用户无需重写既有算法就能把网格优化问题嵌进可微分/机器学习管线，从而做”逆几何处理”。

研究背景

• 领域现状：几何处理里大量任务本质是网格上的优化问题（形变、参数化、测地距离等），社区几十年来积累了很多针对具体问题的高效专用求解器（局部-全局迭代、ADMM、稀疏特征分解等）。
• 核心痛点：如果想把”求解结果”再放进一个更大的优化里（例如反过来求控制点位置，使形变结果逼近目标网格），就需要对”解关于输入的导数”求梯度。现有工具两头不讨好：声明式框架（CVXPYLayers、∇-Prox、Theseus）要求把问题塞进某种规范形式，逼迫用户放弃专用求解器且难以表达非凸问题；命令式做法直接用 PyTorch/JAX 展开（unroll）内层迭代，会随迭代步数堆积中间量、内存爆炸，也无法调用不可微的黑盒高性能求解器。两类工具都很少支持网格上的稀疏、不规则计算。
• 本文 idea：不改写既有算法。让用户只写”一次迭代/一个隐式关系”，系统据此用伴随法在幕后自动生成反向传播，把前向求解器当黑盒，把稀疏性一路保留下来。

方法

整体框架是一个三步式编程模型，对应优化问题 \(\min_{\theta} \ell(g(\varphi(\theta)))\)：用户先把外层变量 \(\theta\)（可以是网格数据，也可以是神经网络参数）通过映射 \(\varphi\) 变成网格上的量 \(x\)；再定义一个”隐式几何层” \(g\) 求解内层几何问题得到 \(y\)；最后用 \(y\) 组一个标量目标 \(\ell\)。整套外层循环写法和普通 PyTorch 完全一样，求导的复杂性被系统藏起来。

flowchart LR
  A["外层参数 θ"] --> B["映射 φ 得到网格量 x"]
  B --> C["隐式几何层 g：解内层问题得 y"]
  C --> D["标量目标 ℓ(y)"]
  D -- "loss.backward()" --> E["伴随法自动反向"]
  E --> A

关键设计：

1. 伴随法作为统一后端。 用户为求解器 \(g\) 额外提供一个隐式关系 \(f(x; y) = 0\)（例如不动点迭代 \(y \leftarrow \tilde f(x;y)\) 就等价于 \(f = \tilde f - y = 0\)）。由隐函数定理，\(\dfrac{\partial y}{\partial x} = -\left(\dfrac{\partial f}{\partial y}\right)^{-1}\dfrac{\partial f}{\partial x}\)。这个公式与前向求解器的实现无关，因此前向可以用任意黑盒（ARPACK、cuDSS、CHOLMOD 等），系统只需要 \(f\) 的偏导。由于损失是标量，反向按伴随法先算 \(\partial \ell / \partial f\)：用 GMRES 迭代反复施加 \(\left(\dfrac{\partial f}{\partial y}\right)^{\top}\) 于 \(\dfrac{\partial \ell}{\partial y}^{\top}\)，再乘上 \(\partial f/\partial x\)。核心洞察是：无论前向多复杂，反向恒为一个稀疏线性求解。
2. 面向网格的 scatter-gather 几何模块。 为了在稠密张量框架里表达稀疏、不规则的网格计算，系统用一套”面层级”表示：把 \(d\) 维面（四面体 / 三角形 / 边 / 顶点）各自编号并赋予规范定向，用少数稠密张量（面-顶点、面-子面、子面-定向）记录拓扑关系。数据在层级间的搬运由 face_index（子面聚到面）和 reduce_on_subface（面散到子面）两个原语完成，底层直接复用 PyTorch 的 scatter-gather。这让网格算子既保留稀疏性，又能与 DLPack、神经网络等生态无缝互通。
3. 对两类高频问题的专用层。 线性系统 \(Ay = b\) 特殊在其反向 Jacobian 恰好还是 \(A\) 本身，可复用符号分解、换用户自定义稀疏求解器；系统手写其伴随实现，允许传入不可微的分解器。特征问题 \(Au = \lambda B u\)（受 \(u^{\top} B u = 1\) 约束）则提供四种求导方式：逐个特征向量解线性系统（作为精度基准）、展开幂迭代、Smirnov-Solomon 截断闭式解、以及基于 QR 不动点关系 \(f = Q((A-\sigma I)^{-1}U) - U = 0\) 的方法，让用户按规模和数量权衡。
4. 命令式而非声明式。 用户直接写一次迭代的 imperative 代码（如 ARAP 的一步局部-全局），用 make_fixed_point_layer / make_adjoint_layer 指定哪些参数是 \(x\)、哪些输出是 \(y\)，系统即挂上隐式反向规则。不需要把非凸的 ARAP 等问题”揉”成某种规范凸形式。

实验结果

论文用四类代表性问题验证：逆平均曲率流（稀疏线性系统）、逆参数化（稀疏特征问题，且与神经场耦合做度量迁移）、逆 ARAP 形变（非凸局部-全局）、逆测地距离（ADMM）。其中最直接的定量对比是逆 ARAP 相对声明式框架 Theseus 的前向/反向耗时与峰值内存（642 顶点球面，25 步优化）：

后端	方法	前向 (s)↓	反向 (s)↓	峰值内存 (GB)↓
CPU	Theseus	43.69	135.38	18.22
CPU	iskra	0.32	0.027	0.67
GPU	Theseus	1.18	2.39	13.76
GPU	iskra	0.095	0.018	0.019

iskra 因能沿用专用局部-全局求解器，前向、反向都快出一到数个数量级，内存低出一个数量级以上；Theseus 在稍大网格上就因内存超过 32GB 而无法运行。其他实验中：逆曲率流上 iskra 比 Theseus 快约两个数量级、比 SparseSolve 快约一个数量级（靠复用符号分解与切换 cuDSS）；逆测地距离上前向与 CVXPYLayers 相当（约 0.3s），反向更快（约 3.67s vs 12.34s），且仅需一步 Loop 细分就能让 CVXPYLayers 内存溢出而 iskra 保持在 1GB 以下。

亮点与局限

• 亮点：
  • “只写一次迭代 + 隐式关系”就能自动获得反向传播，让 30 年积累的专用几何求解器无需重写即可可微化，实现代码往往比声明式方案更短。
  • 用伴随法而非展开迭代，内存不随迭代步数增长，前向可挂任意黑盒/不可微求解器。
  • 面向网格的层级 scatter-gather 表示保留稀疏性，同时通过 PyTorch/DLPack 与机器学习生态互通，支持 CPU/GPU 与神经场耦合。
  • 对特征问题给出四种求导法并系统比较其速度-精度-内存权衡，实用参考价值高。
• 局限：
  • 显式声明不覆盖求解器迭代过程中稀疏结构/网格连通性的动态变化，也不含碰撞检测、接触等物理仿真专属需求。
  • 正确性依赖用户保证 \(f\) 唯一确定 \(y\) 且 \(\partial f/\partial y\) 可逆，否则伴随反向不成立。
  • 反向统一用未做优化的 GMRES（无预条件、内存随迭代线性增长），作者明说未追求峰值性能，稳定性与内存留待未来工作。

延伸思考

这套”隐式微分 + 网格专用稀疏原语”的思路，把可微优化从图像/一般 ML 场景迁到了三角/四面体网格上，恰好补上了可微渲染、神经场之外的一块拼图：外层可以是逆渲染、神经网络预测的度量/目标，内层是经典几何求解器。值得追问的方向包括：如何把它扩展到含动态拓扑与接触的物理仿真（作者列为 future work）；GMRES 反向的预条件与内存优化能带来多大提速；以及在更大规模学习管线里，把 SCP/ARAP/测地这类”几何层”当作可微模块反复调用时的数值稳定性表现。