minidx_core/layers/
gate.rs

1use crate::layers::{Activation, Bias1d, Dense};
2use crate::matmul::MatMulImpl;
3use crate::{Dtype, Float};
4
5/// A gated linear unit, containing dense layers and bias for both the gate and gated connections.
6#[derive(Clone, Debug, Default)]
7pub struct GLU<
8    E: Dtype + Float + MatMulImpl,
9    const I: usize,
10    const O: usize,
11    A: crate::Module<[E; O], Output = [E; O]> + Default = Activation<E>,
12> {
13    gate_connections: Dense<E, I, O>,
14    gate_bias: Bias1d<E, O>,
15
16    sig_connections: Dense<E, I, O>,
17    sig_bias: Bias1d<E, O>,
18
19    activation: A,
20}
21
22impl<E: Dtype + Float + MatMulImpl, const I: usize, const O: usize> GLU<E, I, O, Activation<E>> {
23    pub fn sigmoid() -> Self {
24        Self {
25            activation: Activation::<E>::Sigmoid,
26            ..Self::default()
27        }
28    }
29    pub fn relu() -> Self {
30        Self {
31            activation: Activation::<E>::Relu,
32            ..Self::default()
33        }
34    }
35    pub fn leaky_relu(a: f32) -> Self {
36        Self {
37            activation: Activation::<E>::LeakyRelu(E::from_f32(a).unwrap()),
38            ..Self::default()
39        }
40    }
41    pub fn tanh() -> Self {
42        Self {
43            activation: Activation::<E>::Tanh,
44            ..Self::default()
45        }
46    }
47}
48
49impl<
50        E: Dtype + Float + MatMulImpl,
51        const I: usize,
52        const O: usize,
53        A: crate::Module<[E; O], Output = [E; O]> + Default,
54    > GLU<E, I, O, A>
55{
56    #[doc(hidden)]
57    pub fn connection_params(&self) -> (&[[E; I]; O], &[E; O], &[[E; I]; O], &[E; O]) {
58        (
59            &self.gate_connections.weights,
60            self.gate_bias.bias.raw_grads_ref(),
61            &self.sig_connections.weights,
62            self.sig_bias.bias.raw_grads_ref(),
63        )
64    }
65}
66
67impl<E: Dtype + Float + MatMulImpl, const I: usize, const O: usize>
68    GLU<E, I, O, super::Swish<E, O>>
69{
70    pub fn swish() -> Self {
71        Self::default()
72    }
73}
74
75impl<
76        E: Dtype + Float + MatMulImpl,
77        const I: usize,
78        const O: usize,
79        A: crate::Module<[E; O], Output = [E; O]> + TracedModule<[E; O]> + Default,
80    > crate::Module<[E; I]> for GLU<E, I, O, A>
81{
82    type Output = [E; O];
83
84    fn forward(&self, x: &[E; I]) -> Result<Self::Output, crate::Error> {
85        let gates = self.gate_bias.forward(&self.gate_connections.forward(x)?)?;
86        let mut gates = self.activation.forward(&gates)?;
87        let sigs = self.sig_bias.forward(&self.sig_connections.forward(x)?)?;
88
89        gates
90            .iter_mut()
91            .zip(sigs.into_iter())
92            .for_each(|(o, s)| *o *= s);
93
94        Ok(gates)
95    }
96}
97
98use crate::TracedModule;
99
100impl<
101        E: Dtype + Float + MatMulImpl,
102        const I: usize,
103        const O: usize,
104        A: crate::Module<[E; O], Output = [E; O]> + TracedModule<[E; O]> + Default,
105    > TracedModule<[E; I]> for GLU<E, I, O, A>
106{
107    type Trace = (
108        // gate
109        (
110            <Dense<E, I, O> as TracedModule<[E; I]>>::Trace,
111            <Bias1d<E, O> as TracedModule<[E; O]>>::Trace,
112            <A as TracedModule<[E; O]>>::Trace,
113        ),
114        // signal
115        (
116            <Dense<E, I, O> as TracedModule<[E; I]>>::Trace,
117            <Bias1d<E, O> as TracedModule<[E; O]>>::Trace,
118        ),
119        // outputs for gate and signal
120        ([E; O], [E; O]),
121    );
122
123    fn traced_forward(
124        &self,
125        x: [E; I],
126    ) -> Result<(<Self as crate::Module<[E; I]>>::Output, Self::Trace), crate::Error> {
127        let (gc_out, gc_trace) = self.gate_connections.traced_forward(x.clone())?;
128        let (gb_out, gb_trace) = self.gate_bias.traced_forward(gc_out)?;
129        let (ga_out, ga_trace) = self.activation.traced_forward(gb_out)?;
130
131        let (sc_out, sc_trace) = self.sig_connections.traced_forward(x)?;
132        let (sb_out, sb_trace) = self.sig_bias.traced_forward(sc_out)?;
133
134        let mut out = ga_out.clone();
135        out.iter_mut()
136            .zip(sb_out.iter())
137            .for_each(|(o, s)| *o *= *s);
138
139        Ok((
140            out,
141            (
142                (gc_trace, gb_trace, ga_trace),
143                (sc_trace, sb_trace),
144                (ga_out, sb_out),
145            ),
146        ))
147    }
148}
149
150use crate::BackpropModule;
151
152impl<
153        E: Dtype + Float + MatMulImpl,
154        const I: usize,
155        const O: usize,
156        A: crate::Module<[E; O], Output = [E; O]>
157            + TracedModule<[E; O]>
158            + BackpropModule<[E; O]>
159            + Default,
160    > crate::BackpropModule<[E; I]> for GLU<E, I, O, A>
161{
162    type SelfGrads = (
163        // gate
164        (
165            <Dense<E, I, O> as BackpropModule<[E; I]>>::SelfGrads,
166            <Bias1d<E, O> as BackpropModule<[E; O]>>::SelfGrads,
167            <A as BackpropModule<[E; O]>>::SelfGrads,
168        ),
169        // signal
170        (
171            <Dense<E, I, O> as BackpropModule<[E; I]>>::SelfGrads,
172            <Bias1d<E, O> as BackpropModule<[E; O]>>::SelfGrads,
173        ),
174    );
175
176    fn backprop(
177        &self,
178        trace: &<Self as crate::TracedModule<[E; I]>>::Trace,
179        grads_wrt_output: <Self as crate::Module<[E; I]>>::Output,
180    ) -> ([E; I], Self::SelfGrads) {
181        // grads for the signals side of the chain is our output grads * the output of the gate
182        let mut sig_grads_wrt_output = grads_wrt_output.clone();
183        sig_grads_wrt_output
184            .iter_mut()
185            .zip(trace.2 .0)
186            .for_each(|(g, o)| *g *= o);
187        // backprop for the signals chain
188        let (sig_grads, sb_grads) = self.sig_bias.backprop(&trace.1 .1, sig_grads_wrt_output);
189        let (sig_grads, sc_grads) = self.sig_connections.backprop(&trace.1 .0, sig_grads);
190
191        // grads for the gate side of the chain is our output grads * the output of the signals
192        let mut gate_grads_wrt_output = grads_wrt_output;
193        gate_grads_wrt_output
194            .iter_mut()
195            .zip(trace.2 .1)
196            .for_each(|(g, o)| *g *= o);
197
198        // backprop for the signals chain
199        let (gate_grads, ga_grads) = self.activation.backprop(&trace.0 .2, gate_grads_wrt_output);
200        let (gate_grads, gb_grads) = self.gate_bias.backprop(&trace.0 .1, gate_grads);
201        let (gate_grads, gc_grads) = self.gate_connections.backprop(&trace.0 .0, gate_grads);
202
203        // output derivatives is just the sum of the partial derivatives from each chain
204        // TODO: not sure about this
205        let mut out = sig_grads;
206        out.iter_mut()
207            .zip(gate_grads.into_iter())
208            .for_each(|(o, x)| *o += x);
209
210        (out, ((gc_grads, gb_grads, ga_grads), (sc_grads, sb_grads)))
211    }
212
213    fn update(
214        &mut self,
215        applyer: &mut impl crate::optimizers::GradApplyer,
216        updates: Self::SelfGrads,
217    ) -> Result<(), crate::Error> {
218        self.gate_connections.update(applyer, updates.0 .0)?;
219        self.gate_bias.update(applyer, updates.0 .1)?;
220        self.activation.update(applyer, updates.0 .2)?;
221
222        self.sig_connections.update(applyer, updates.1 .0)?;
223        self.sig_bias.update(applyer, updates.1 .1)?;
224
225        Ok(())
226    }
227}
228
229impl<
230        E: Dtype + Float + MatMulImpl,
231        const I: usize,
232        const O: usize,
233        A: crate::Module<[E; O], Output = [E; O]> + Default + crate::LoadableModule,
234    > crate::LoadableModule for GLU<E, I, O, A>
235{
236    fn save(
237        &self,
238        path: String,
239        dict: &mut std::collections::HashMap<String, Vec<f64>>,
240    ) -> Result<(), crate::LoadSaveError> {
241        self.gate_connections
242            .save(path.clone() + ".gate_connections", dict)?;
243        self.gate_bias.save(path.clone() + ".gate_bias", dict)?;
244        self.activation.save(path.clone() + ".activation", dict)?;
245
246        self.sig_connections
247            .save(path.clone() + ".sig_connections", dict)?;
248        self.sig_bias.save(path + ".sig_bias", dict)?;
249
250        Ok(())
251    }
252
253    fn load(
254        &mut self,
255        path: String,
256        dict: &std::collections::HashMap<String, Vec<f64>>,
257    ) -> Result<(), crate::LoadSaveError> {
258        self.gate_connections
259            .load(path.clone() + ".gate_connections", dict)?;
260        self.gate_bias.load(path.clone() + ".gate_bias", dict)?;
261        self.activation.load(path.clone() + ".activation", dict)?;
262
263        self.sig_connections
264            .load(path.clone() + ".sig_connections", dict)?;
265        self.sig_bias.load(path + ".sig_bias", dict)?;
266
267        Ok(())
268    }
269}
270
271impl<
272        E: Dtype + Float + MatMulImpl,
273        const I: usize,
274        const O: usize,
275        A: crate::Module<[E; O], Output = [E; O]> + Default + crate::ResetParams,
276    > crate::ResetParams for GLU<E, I, O, A>
277{
278    fn rand_params<RNG: rand::Rng>(
279        &mut self,
280        rng: &mut RNG,
281        scale: f32,
282    ) -> Result<(), crate::Error> {
283        self.gate_connections.rand_params(rng, scale)?;
284        self.gate_bias.rand_params(rng, scale)?;
285        self.sig_connections.rand_params(rng, scale)?;
286        self.sig_bias.rand_params(rng, scale)?;
287        self.activation.rand_params(rng, scale)?;
288        Ok(())
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295    use crate::Module;
296
297    #[test]
298    fn test_forward() {
299        let mut g = GLU::<f32, 1, 1>::default();
300        g.gate_connections.weights[0] = [1.0];
301        g.sig_connections.weights[0] = [1.0];
302
303        assert_eq!(g.forward(&[1.0]), Ok([1.0]));
304
305        g.activation = Activation::Sigmoid;
306        // o = 0.5 * sigmoid(0.5) - which itself is 0.62...
307        let o = g.forward(&[0.5]).unwrap()[0];
308        assert!(o > 0.3);
309        assert!(o < 0.35);
310
311        let mut g = GLU::<f32, 2, 1>::default();
312        g.gate_connections.weights[0] = [1.0, 0.0];
313        g.sig_connections.weights[0] = [0.0, 1.0];
314
315        assert_eq!(g.forward(&[1.0, 1.0]), Ok([1.0]));
316        assert_eq!(g.forward(&[1.0, 0.5]), Ok([0.5]));
317        assert_eq!(g.forward(&[0.2, 1.0]), Ok([0.2]));
318        assert_eq!(g.forward(&[-1.0, 1.0]), Ok([0.0])); // relu on gate
319
320        assert_eq!(g.traced_forward([1.0, 0.5]).unwrap().0, [0.5]);
321        assert_eq!(g.traced_forward([-1.0, 1.0]).unwrap().0, [0.0]); // relu on gate
322    }
323
324    #[test]
325    fn test_backward_simple() {
326        let mut g = GLU::<f32, 2, 1>::default();
327        g.gate_connections.weights = [[2.0, 1.0]];
328        g.gate_bias.bias.raw_grads_mut()[0] = -2.0;
329        g.sig_connections.weights = [[1.0, -1.0]];
330        g.sig_bias.bias.raw_grads_mut()[0] = 1.0;
331
332        let (out, trace) = g.traced_forward([1.0, 2.0]).unwrap();
333        assert_eq!(out, [0.0]);
334
335        // trace.2 = (gate_out, sig_out)
336        assert_eq!(trace.2, ([2.0], [0.0]));
337
338        // assume MSE loss, want=1 so loss=0.5 and dL/dY = -1.
339        let (_out_grads, grads) = g.backprop(&trace, [-1.0]);
340        assert_eq!(out, [0.0]); // no gradient backwards as the gate was inactive
341
342        assert_eq!(
343            (grads.0 .0, grads.0 .1.raw_grads(), grads.0 .2),
344            ([[0.0, 0.0]], [0.0], ())
345        ); // no gradient for gate as sig was inactive
346        assert_eq!(
347            (grads.1 .0, grads.1 .1.raw_grads()),
348            ([[-2.0, -4.0]], [-2.0])
349        ); // activation was non-zero so gradients for gate
350    }
351}
minidx_core/layers/gate.rs

minidx_core/layers/
gate.rs
