deep_delta_learn/
nn.rs

1//! Neural network modules for Deep Delta Learning
2//!
3//! Conventions (important in Burn):
4//! - `nn::Linear` is rank-preserving: it applies to the last dimension and keeps all leading dims.
5//!   If you feed `[B, D, V]` into Linear(D -> H), you'll get `[B, D, H]` (still 3D).
6//! - Therefore, branches that *must* return rank-2 tensors (e.g. k: [B, D], beta: [B, 1], v: [B, V])
7//!   should **explicitly pool/flatten** before any Linear layers.
8//!
9//! This module defines:
10//! - `DeltaResidual`: a full Delta-Res block using internal generator branches.
11//! - `DeltaResidualExternal`: a variant that takes externally-provided (k, beta, v), useful for tauformer.
12
13use burn::prelude::*;
14
15use crate::branches::{DeltaBranches, DeltaBranchesConfig};
16use crate::delta::delta_update;
17
18/// Configuration for the DeltaResidual block.
19#[derive(Config, Debug)]
20pub struct DeltaResidualConfig {
21    /// Feature dimension d
22    pub d_model: usize,
23    /// Value dimension dv
24    pub dv: usize,
25    /// Hidden dimension used in generator MLPs
26    pub d_hidden: usize,
27    /// Numerical epsilon used in delta_update and k-normalization
28    #[config(default = "1e-8")]
29    pub eps: f32,
30}
31
32/// Delta Residual block (Delta-Res).
33///
34/// Implements Eq. (2.5):
35/// X_{l+1} = X_l + β(X_l) k(X_l) (v(X_l)^T - k(X_l)^T X_l)
36#[derive(Module, Debug)]
37pub struct DeltaResidual<B: Backend> {
38    branches: DeltaBranches<B>,
39    eps: f32,
40}
41
42impl DeltaResidualConfig {
43    pub fn init<B: Backend>(&self, device: &B::Device) -> DeltaResidual<B> {
44        let branches_cfg = DeltaBranchesConfig::new(self.d_model, self.dv, self.d_hidden);
45        let branches_cfg = branches_cfg.with_eps(self.eps);
46
47        DeltaResidual {
48            branches: branches_cfg.init(device),
49            eps: self.eps,
50        }
51    }
52}
53
54impl<B: Backend> DeltaResidual<B> {
55    /// Forward for matrix state.
56    ///
57    /// Signature convention:
58    /// - input:  `x` is always `[B, D, V]` (`Tensor<B, 3>`)
59    /// - output: `y` is always `[B, D, V]` (`Tensor<B, 3>`)
60    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
61        let (k, beta, v) = self.branches.forward(&x);
62        delta_update(x, k, v, beta, self.eps)
63    }
64
65    /// Like `forward`, but also returns the generated parameters (useful for debugging/tests).
66    pub fn forward_with_params(
67        &self,
68        x: Tensor<B, 3>,
69    ) -> (Tensor<B, 3>, Tensor<B, 2>, Tensor<B, 2>, Tensor<B, 2>) {
70        let (k, beta, v) = self.branches.forward(&x);
71        let y = delta_update(x, k.clone(), v.clone(), beta.clone(), self.eps);
72        (y, k, beta, v)
73    }
74}
75
76/// Delta residual that takes external (k, beta, v).
77#[derive(Module, Debug)]
78pub struct DeltaResidualExternal<B: Backend> {
79    eps: f32,
80    _backend: core::marker::PhantomData<B>, // Required if no other B-related fields exist
81}
82
83impl<B: Backend> DeltaResidualExternal<B> {
84    pub fn new(eps: f32) -> Self {
85        Self {
86            eps,
87            _backend: core::marker::PhantomData,
88        }
89    }
90
91    /// Forward with external parameters.
92    ///
93    /// Shapes:
94    /// - x:    [B, D, V]
95    /// - k:    [B, D]
96    /// - v:    [B, V]
97    /// - beta: [B, 1]
98    pub fn forward(
99        &self,
100        x: Tensor<B, 3>,
101        k: Tensor<B, 2>,
102        v: Tensor<B, 2>,
103        beta: Tensor<B, 2>,
104    ) -> Tensor<B, 3> {
105        delta_update(x, k, v, beta, self.eps)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    pub type TestBackend = burn_ndarray::NdArray<f32>;
114    use burn::tensor::Distribution;
115
116    #[test]
117    fn test_delta_residual_shapes() {
118        let device = Default::default();
119
120        let cfg = DeltaResidualConfig::new(8, 4, 16);
121        let model = cfg.init::<TestBackend>(&device);
122
123        let x =
124            Tensor::<TestBackend, 3>::random([2, 8, 4], Distribution::Uniform(-1.0, 1.0), &device);
125
126        let (y, k, beta, v) = model.forward_with_params(x);
127        assert_eq!(y.dims(), [2, 8, 4]);
128        assert_eq!(k.dims(), [2, 8]);
129        assert_eq!(beta.dims(), [2, 1]);
130        assert_eq!(v.dims(), [2, 4]);
131    }
132
133    #[test]
134    fn test_delta_residual_external_shapes() {
135        let device = Default::default();
136
137        let block = DeltaResidualExternal::new(1e-8);
138
139        let x =
140            Tensor::<TestBackend, 3>::random([2, 8, 4], Distribution::Uniform(-1.0, 1.0), &device);
141        let k = Tensor::<TestBackend, 2>::random([2, 8], Distribution::Uniform(-1.0, 1.0), &device);
142        let v = Tensor::<TestBackend, 2>::random([2, 4], Distribution::Uniform(-1.0, 1.0), &device);
143        let beta =
144            Tensor::<TestBackend, 2>::random([2, 1], Distribution::Uniform(0.0, 2.0), &device);
145
146        let y = block.forward(x, k, v, beta);
147        assert_eq!(y.dims(), [2, 8, 4]);
148    }
149}
deep_delta_learn/nn.rs

deep_delta_learn/
nn.rs
