deep_delta_learn/
branches.rs

1//! Generator branches for Deep Delta Learning parameters.
2//!
3//! Burn tensor shape conventions (critical):
4//! - Tensor rank is part of the type: `Tensor<B, const D: usize, K>`.
5//! - Many operations in Burn are rank-preserving unless you explicitly request a new rank
6//!   using APIs like `flatten::<D2>(..)`, `reshape::<D2>(..)`, `squeeze::<D2>()`.
7//! - `nn::Linear` is rank-preserving: it transforms the last dimension and keeps all leading dims.
8//!
9//! Therefore, to ensure branches return rank-2 tensors:
10//! - Always convert the input state `x: Tensor<B, 3>` into an explicit rank-2 tensor
11//!   before passing to any Linear.
12//!
13//! This module follows Appendix A from the paper (arXiv:2601.00417v1). [file:1]
14
15use burn::nn::{Linear, LinearConfig};
16use burn::prelude::*;
17
18/// Convert X [B, D, V] into a pooled representation [B, D] by averaging over V.
19///
20/// Implementation detail:
21/// - We compute a mean over dim=2 but do not assume the rank changes.
22/// - If mean keeps dims, we squeeze the singleton dimension.
23///
24/// Returned tensor is always `Tensor<B, 2>`.
25#[inline]
26pub fn pool_v_to_bd<B: Backend>(x: &Tensor<B, 3>) -> Tensor<B, 2> {
27    // Try to keep semantics robust: mean_dim might return [B, D, 1] or [B, D] depending on version.
28    // We force rank-2 by squeezing dim=2 if it's a singleton.
29    let pooled = x.clone().mean_dim(2);
30
31    // If pooled is already rank-2 in your Burn version, this line won't compile.
32    // So we implement using `flatten` in a way that always compiles:
33    // Strategy: compute sum over V explicitly using elementwise operations is not available here.
34    // Instead, we rely on the fact that in current Burn, `mean_dim` returns same rank (keeps dim).
35    // Then we squeeze to rank-2.
36    let pooled: Tensor<B, 2> = pooled.squeeze::<2>(2);
37    pooled
38}
39
40/// Convert X [B, D, V] into a pooled representation [B, V] by averaging over D.
41#[inline]
42pub fn pool_d_to_bv<B: Backend>(x: &Tensor<B, 3>) -> Tensor<B, 2> {
43    let pooled = x.clone().mean_dim(1);
44    let pooled: Tensor<B, 2> = pooled.squeeze::<2>(1);
45    pooled
46}
47
48/// L2-normalize vectors along the last dimension for a rank-2 tensor [B, D].
49#[inline]
50pub fn l2_normalize_bd<B: Backend>(x: Tensor<B, 2>, eps: f32) -> Tensor<B, 2> {
51    // norm: [B]
52    let norm: Tensor<B, 2> = x.clone().powf_scalar(2.0).sum_dim(1).sqrt().add_scalar(eps);
53    // [B] -> [B, 1]
54    let norm: Tensor<B, 2> = norm.unsqueeze::<2>();
55    x.div(norm)
56}
57
58// ---------------------------
59// K-Branch
60// ---------------------------
61
62#[derive(Config, Debug)]
63pub struct KBranchConfig {
64    pub d_model: usize,
65    pub d_hidden: usize,
66    #[config(default = "1e-8")]
67    pub eps: f32,
68}
69
70/// k(X) generator.
71///
72/// k̃ = MLP(Pool(X)), k = k̃ / (||k̃|| + eps) [file:1]
73#[derive(Module, Debug)]
74pub struct KBranch<B: Backend> {
75    linear1: Linear<B>,
76    linear2: Linear<B>,
77    eps: f32,
78}
79
80impl KBranchConfig {
81    pub fn init<B: Backend>(&self, device: &B::Device) -> KBranch<B> {
82        KBranch {
83            linear1: LinearConfig::new(self.d_model, self.d_hidden).init(device),
84            linear2: LinearConfig::new(self.d_hidden, self.d_model).init(device),
85            eps: self.eps,
86        }
87    }
88}
89
90impl<B: Backend> KBranch<B> {
91    /// X: [B, D, V] -> k: [B, D]
92    pub fn forward(&self, x: &Tensor<B, 3>) -> Tensor<B, 2> {
93        let pooled: Tensor<B, 2> = pool_v_to_bd(x); // [B, D]
94        let h: Tensor<B, 2> = self.linear1.forward(pooled);
95        let h: Tensor<B, 2> = burn::tensor::activation::gelu(h);
96        let k: Tensor<B, 2> = self.linear2.forward(h);
97        l2_normalize_bd(k, self.eps)
98    }
99}
100
101// ---------------------------
102// Beta-Branch
103// ---------------------------
104
105#[derive(Config, Debug)]
106pub struct BetaBranchConfig {
107    pub d_model: usize,
108    pub d_hidden: usize,
109}
110
111/// β(X) generator.
112///
113/// β(X) = 2 * sigmoid( ... ) in range [0,2]. [file:1]
114#[derive(Module, Debug)]
115pub struct BetaBranch<B: Backend> {
116    w_in: Linear<B>,
117    w_out: Linear<B>,
118}
119
120impl BetaBranchConfig {
121    pub fn init<B: Backend>(&self, device: &B::Device) -> BetaBranch<B> {
122        BetaBranch {
123            w_in: LinearConfig::new(self.d_model, self.d_hidden).init(device),
124            w_out: LinearConfig::new(self.d_hidden, 1).init(device),
125        }
126    }
127}
128
129impl<B: Backend> BetaBranch<B> {
130    /// X: [B, D, V] -> beta: [B, 1]
131    pub fn forward(&self, x: &Tensor<B, 3>) -> Tensor<B, 2> {
132        use burn::tensor::activation::sigmoid;
133        let pooled: Tensor<B, 2> = pool_v_to_bd(x); // [B, D]
134        let h: Tensor<B, 2> = self.w_in.forward(pooled).tanh();
135        let logit: Tensor<B, 2> = self.w_out.forward(h); // [B, 1]
136        sigmoid(logit).mul_scalar(2.0)
137    }
138}
139
140// ---------------------------
141// V-Branch
142// ---------------------------
143
144#[derive(Config, Debug)]
145pub struct VBranchConfig {
146    pub d_model: usize,
147    pub dv: usize,
148    pub d_hidden: usize,
149}
150
151/// v(X) generator.
152///
153/// The paper treats v(X) as a flexible content-update branch. [file:1]
154///
155/// Default here:
156/// - pool over V to get [B, D]
157/// - MLP D -> hidden -> dv
158#[derive(Module, Debug)]
159pub struct VBranch<B: Backend> {
160    linear1: Linear<B>,
161    linear2: Linear<B>,
162}
163
164impl VBranchConfig {
165    pub fn init<B: Backend>(&self, device: &B::Device) -> VBranch<B> {
166        VBranch {
167            linear1: LinearConfig::new(self.d_model, self.d_hidden).init(device),
168            linear2: LinearConfig::new(self.d_hidden, self.dv).init(device),
169        }
170    }
171}
172
173impl<B: Backend> VBranch<B> {
174    /// X: [B, D, V] -> v: [B, V]
175    pub fn forward(&self, x: &Tensor<B, 3>) -> Tensor<B, 2> {
176        let pooled: Tensor<B, 2> = pool_v_to_bd(x); // [B, D]
177        let h: Tensor<B, 2> = self.linear1.forward(pooled);
178        let h: Tensor<B, 2> = burn::tensor::activation::gelu(h);
179        let v: Tensor<B, 2> = self.linear2.forward(h);
180        v
181    }
182}
183
184// ---------------------------
185// Bundle
186// ---------------------------
187
188#[derive(Config, Debug)]
189pub struct DeltaBranchesConfig {
190    pub d_model: usize,
191    pub dv: usize,
192    pub d_hidden: usize,
193    #[config(default = "1e-8")]
194    pub eps: f32,
195}
196
197#[derive(Module, Debug)]
198pub struct DeltaBranches<B: Backend> {
199    pub k_branch: KBranch<B>,
200    pub beta_branch: BetaBranch<B>,
201    pub v_branch: VBranch<B>,
202}
203
204impl DeltaBranchesConfig {
205    pub fn init<B: Backend>(&self, device: &B::Device) -> DeltaBranches<B> {
206        DeltaBranches {
207            k_branch: KBranchConfig::new(self.d_model, self.d_hidden)
208                .with_eps(self.eps)
209                .init(device),
210            beta_branch: BetaBranchConfig::new(self.d_model, self.d_hidden).init(device),
211            v_branch: VBranchConfig::new(self.d_model, self.dv, self.d_hidden).init(device),
212        }
213    }
214}
215
216impl<B: Backend> DeltaBranches<B> {
217    /// X: [B, D, V] -> (k: [B, D], beta: [B, 1], v: [B, V])
218    pub fn forward(&self, x: &Tensor<B, 3>) -> (Tensor<B, 2>, Tensor<B, 2>, Tensor<B, 2>) {
219        let k = self.k_branch.forward(x);
220        let beta = self.beta_branch.forward(x);
221        let v = self.v_branch.forward(x);
222        (k, beta, v)
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    use crate::backend::AutoBackend;
231    type TestBackend = AutoBackend;
232    use burn::tensor::Distribution;
233
234    #[test]
235    fn test_pool_v_to_bd_shape() {
236        let device = Default::default();
237        let x =
238            Tensor::<TestBackend, 3>::random([2, 8, 4], Distribution::Uniform(-1.0, 1.0), &device);
239        let pooled = pool_v_to_bd(&x);
240        assert_eq!(pooled.dims(), [2, 8]);
241    }
242
243    #[test]
244    fn test_k_branch_shape_and_norm() {
245        let device = Default::default();
246        let cfg = KBranchConfig::new(8, 16);
247        let branch = cfg.init::<TestBackend>(&device);
248
249        let x =
250            Tensor::<TestBackend, 3>::random([2, 8, 4], Distribution::Uniform(-1.0, 1.0), &device);
251        let k = branch.forward(&x);
252        assert_eq!(k.dims(), [2, 8]);
253
254        let norms: Tensor<TestBackend, 2> = k.clone().powf_scalar(2.0).sum_dim(1).sqrt();
255        let norms: Vec<f32> = norms.into_data().to_vec().unwrap();
256        for n in norms {
257            assert!((n - 1.0).abs() < 0.05);
258        }
259    }
260
261    #[test]
262    fn test_beta_branch_range() {
263        let device = Default::default();
264        let cfg = BetaBranchConfig::new(8, 16);
265        let branch = cfg.init::<TestBackend>(&device);
266
267        let x =
268            Tensor::<TestBackend, 3>::random([2, 8, 4], Distribution::Uniform(-1.0, 1.0), &device);
269        let beta = branch.forward(&x);
270        assert_eq!(beta.dims(), [2, 1]);
271
272        let vals: Vec<f32> = beta.into_data().to_vec().unwrap();
273        for b in vals {
274            assert!(b >= 0.0 && b <= 2.0);
275        }
276    }
277
278    #[test]
279    fn test_v_branch_shape() {
280        let device = Default::default();
281        let cfg = VBranchConfig::new(8, 4, 16);
282        let branch = cfg.init::<TestBackend>(&device);
283
284        let x =
285            Tensor::<TestBackend, 3>::random([2, 8, 4], Distribution::Uniform(-1.0, 1.0), &device);
286        let v = branch.forward(&x);
287        assert_eq!(v.dims(), [2, 4]);
288    }
289
290    #[test]
291    fn test_bundle_shapes() {
292        let device = Default::default();
293        let cfg = DeltaBranchesConfig::new(8, 4, 16);
294        let branches = cfg.init::<TestBackend>(&device);
295
296        let x =
297            Tensor::<TestBackend, 3>::random([2, 8, 4], Distribution::Uniform(-1.0, 1.0), &device);
298        let (k, beta, v) = branches.forward(&x);
299
300        assert_eq!(k.dims(), [2, 8]);
301        assert_eq!(beta.dims(), [2, 1]);
302        assert_eq!(v.dims(), [2, 4]);
303    }
304}
deep_delta_learn/branches.rs

deep_delta_learn/
branches.rs
