burn_mamba/mamba3/quat_scan/
backward.rs

1//! # Custom autodiff node for the quaternion cumulative-product scan
2//!
3//! Implements [`Mamba3QuatScanBackendExt`](super::quat_scan::Mamba3QuatScanBackendExt)
4//! for `Autodiff<B>` via a single Burn [`Backward`] node. The forward keeps only
5//! its two leaf inputs (the per-step quaternions `q` and the carry `init`);
6//! backprop recomputes the prefix product and evaluates the exact quaternion VJP
7//! of the cumulative product, so the `O(log seq)` scan intermediates are never
8//! retained.
9//!
10//! ## Gradient math
11//!
12//! The scan is `cum[t] = qₜ ⊗ cum[t-1]`, `cum[-1] = init` (so `cum[t] = Pₜ ⊗ init`
13//! with the prefix product `Pₜ = qₜ ⊗ ⋯ ⊗ q₀`). Quaternion multiplication is
14//! bilinear, and for the real inner product the per-factor VJPs are exact for
15//! **all** quaternions:
16//!
17//! ```text
18//!   o = a ⊗ b ,  cotangent ḡ   ⟹   grad_a = ḡ ⊗ conj(b) ,  grad_b = conj(a) ⊗ ḡ .
19//! ```
20//!
21//! Let `G[t]` be the total cotangent reaching `cum[t]` (the upstream `d_cum[t]`
22//! plus what flows back from `cum[t+1]`). The reverse recurrence
23//! `G[t] = d_cum[t] + conj(qₜ₊₁) ⊗ G[t+1]` telescopes, and **for unit
24//! quaternions** (which the rotation always produces) collapses to a
25//! parallel-friendly closed form:
26//!
27//! ```text
28//!   S[t]   = Σ_{s≥t} conj(Pₛ) ⊗ d_cum[s]        (suffix-sum = reverse cumsum)
29//!   G[t]   = Pₜ ⊗ S[t]
30//!   d_q[t] = G[t] ⊗ conj(cum[t-1])              (cum[-1] = init)
31//!   d_init = S[0]
32//! ```
33//!
34//! `final_carry = cum[:, −1]` is a high-level autodiff slice (in
35//! [`quat_cumprod_recalculated`](super::quat_scan::quat_cumprod_recalculated)),
36//! so its gradient is already folded into `d_cum` at the last position before
37//! this node runs — the node has a single output.
38
39#![allow(non_snake_case)]
40
41use super::quat_scan::{Mamba3QuatScanBackendExt, Quat, quat_prefix_product_soa};
42use crate::utils::fprim::F;
43use burn::backend::autodiff::{
44    Autodiff,
45    checkpoint::{base::Checkpointer, strategy::CheckpointStrategy},
46    grads::Gradients,
47    ops::{Backward, Ops, OpsKind},
48};
49use burn::backend::tensor::FloatTensor;
50use burn::backend::{Backend, BackendTypes};
51
52/// Recompute-based gradient of the quaternion cumulative product.
53///
54/// Given the saved inputs (`q`, `init`) and the upstream cotangent `d_cum`,
55/// recomputes the prefix product `P` and returns `(d_q, d_init)` via the
56/// closed-form unit-quaternion VJP documented in the module header. All ops are
57/// parallel (a prefix product + a reverse-cumsum); there is no token loop.
58fn quat_scan_backward<B: Backend>(
59    q_bshj4: F<B, 5>,
60    init_bhj4: F<B, 4>,
61    d_cum_bshj4: F<B, 5>,
62) -> (F<B, 5>, F<B, 4>) {
63    // Unpack to struct-of-arrays once; all the heavy math is narrow/cat-free.
64    let q = Quat::from_rank5(q_bshj4);
65    let init = Quat::from_rank5(init_bhj4.unsqueeze_dim::<5>(1)); // [b,1,h,j] components
66    let d_cum = Quat::from_rank5(d_cum_bshj4);
67    let sequence = q.w.dims()[1];
68
69    // Pₜ = qₜ ⊗ … ⊗ q₀ (prefix product, no carry) — recomputed, not stored.
70    let p = quat_prefix_product_soa::<B>(q);
71
72    // cum[t] = Pₜ ⊗ init (broadcast over the sequence axis).
73    let cum = p.clone().mul(init.clone());
74
75    // S[t] = Σ_{s≥t} conj(Pₛ) ⊗ d_cum[s] — suffix-sum (reverse cumsum) per component.
76    let s = p.clone().conj().mul(d_cum).reverse_cumsum_seq();
77
78    // G[t] = Pₜ ⊗ S[t].
79    let g = p.mul(s.clone());
80
81    // cum_prev[t] = cum[t-1]; cum_prev[0] = init (prepend init, drop the last).
82    let cum_prev = cum.shift_prepend(init, sequence);
83
84    // d_q[t] = G[t] ⊗ conj(cum_prev[t]);  d_init = S[0].
85    let d_q_bshj4 = g.mul(cum_prev.conj()).pack();
86    let d_init_bhj4 = s.pack().narrow(1, 0, 1).squeeze_dim::<4>(1);
87
88    (d_q_bshj4, d_init_bhj4)
89}
90
91impl<B: Backend + Mamba3QuatScanBackendExt, C: CheckpointStrategy> Mamba3QuatScanBackendExt
92    for Autodiff<B, C>
93{
94    fn quat_cumprod(q_bshj4: FloatTensor<Self>, init_bhj4: FloatTensor<Self>) -> FloatTensor<Self> {
95        // ── Backward node definition ─────────────────────────────────────────
96        #[derive(Debug)]
97        struct QuatScanBackward;
98
99        #[derive(Clone, Debug)]
100        struct State<B: Backend> {
101            // Saved forward inputs (only the leaves — no scan intermediates).
102            q_bshj4: <B as BackendTypes>::FloatTensorPrimitive,
103            init_bhj4: <B as BackendTypes>::FloatTensorPrimitive,
104            shape_q_bshj4: [usize; 5],
105            shape_init_bhj4: [usize; 4],
106        }
107
108        impl<B: Backend + Mamba3QuatScanBackendExt> Backward<B, 2> for QuatScanBackward {
109            type State = State<B>;
110
111            fn backward(
112                self,
113                ops: Ops<Self::State, 2>,
114                grads: &mut Gradients,
115                _checkpointer: &mut Checkpointer,
116            ) {
117                let [node_q, node_init] = ops.parents;
118
119                // Upstream gradient of `cum` (already includes the final_carry
120                // slice's contribution, scattered in at the last position).
121                let d_cum: <B as BackendTypes>::FloatTensorPrimitive =
122                    grads.consume::<B>(&ops.node);
123
124                let State {
125                    q_bshj4,
126                    init_bhj4,
127                    shape_q_bshj4,
128                    shape_init_bhj4,
129                } = ops.state;
130
131                let q = F::<B, 5>::new(q_bshj4).reshape(shape_q_bshj4);
132                let init = F::<B, 4>::new(init_bhj4).reshape(shape_init_bhj4);
133                // `cum` has the same shape as `q`, hence so does its gradient.
134                let d_cum = F::<B, 5>::new(d_cum).reshape(shape_q_bshj4);
135
136                let (d_q, d_init) = quat_scan_backward::<B>(q, init, d_cum);
137
138                if let Some(n) = node_q {
139                    grads.register::<B>(n.id, d_q.inner());
140                }
141                if let Some(n) = node_init {
142                    grads.register::<B>(n.id, d_init.inner());
143                }
144            }
145        }
146
147        // ── Shape extraction ───────────────────────────────────────────────
148        use burn::backend::TensorMetadata;
149        let shape_q_bshj4: [usize; 5] = q_bshj4.primitive.shape().dims();
150        let shape_init_bhj4: [usize; 4] = init_bhj4.primitive.shape().dims();
151
152        // ── Register backward / run forward ─────────────────────────────────
153        match QuatScanBackward
154            .prepare::<C>([q_bshj4.node.clone(), init_bhj4.node.clone()])
155            .compute_bound()
156            .stateful()
157        {
158            OpsKind::Tracked(prep) => {
159                let prim_cum =
160                    B::quat_cumprod(q_bshj4.primitive.clone(), init_bhj4.primitive.clone());
161                let state = State {
162                    q_bshj4: q_bshj4.primitive.clone(),
163                    init_bhj4: init_bhj4.primitive.clone(),
164                    shape_q_bshj4,
165                    shape_init_bhj4,
166                };
167                prep.finish(state, prim_cum)
168            }
169            OpsKind::UnTracked(prep) => {
170                let prim_cum = B::quat_cumprod(q_bshj4.primitive, init_bhj4.primitive);
171                prep.finish(prim_cum)
172            }
173        }
174    }
175}
burn_mamba/mamba3/quat_scan/backward.rs

burn_mamba/mamba3/quat_scan/
backward.rs
