ICML2026

d2^2p: Structured Soft Attention Is All You Need

Casey Mogilevsky, Kimberly Liang

摘要

Classical dynamic programming algorithms such as Smith-Waterman, edit distance, and constituency parsing solve structured combinatorial problems using hard constraints. Soft relaxations replace hard max\max or min\min operators with differentiable probabilistic models whose gradients correspond to alignment or parsing marginals, but existing approaches typically treat the DP algorithm itself as fixed, relying on hand-tuned gap penalties, edit costs, span penalties, and temperatures. We show how to learn these parameters directly from data: the marginals these algorithms produce are themselves first-order derivatives, so learning them is inherently a second-order problem. We derive efficient Hessian-vector products and cross-Jacobians for twelve dynamic programming algorithms spanning alignment, edit distance, and parsing; both derivative families admit closed-form covariance expressions under the induced Gibbs distribution. We implement these operators as fused CUDA kernels in d2^2p, achieving 100100--20,000×20{,}000\times speedups over standard PyTorch and making end-to-end parameter learning practical at modern scales. We then demonstrate that learning these parameters is critical in practice. In protein structure alignment, freezing gap penalties collapses performance from 0.740.74 to 0.320.32 F1F_1 (and to 0.130.13 at lower encoder capacity), while jointly learning them recovers biologically meaningful gap regimes and reaches 0.750.75 F1F_1 and 0.4450.445 lDDT, 9191% of the TM-align lDDT ceiling. The same machinery transfers to constituency parsing, where a structured CKY CRF matches dense per-span supervision to within 0.0030.003 F1F_1 on English with no direct parse-tree structural supervision.