Initial commit

src/transforms.rs

@@ -0,0 +1,315 @@
+use crate::float::Float;
+use crate::linear_algebra::{ColumnEfficientMatrix, Vector};
+use rand_distr::Normal;
+
+pub fn l2_error(xs: &[Float], ys: &[Float]) -> Float {
+    // err-squared
+    let mut result = 0.0;
+    for (x, y) in xs.iter().zip(ys) {
+        let s = x - y;
+        result += s * s;
+    }
+    0.5 * result
+}
+
+pub fn gradient_l2_error_mut(xs: &[Float], ys: &[Float], out: &mut [Float]) {
+    for (i, (x, y)) in xs.iter().zip(ys).enumerate() {
+        out[i] = x - y;
+    }
+}
+
+fn sigmoid(x: Float) -> Float {
+    1.0 / (1.0 + (-x).exp())
+}
+
+fn activation_of_sigmoid_potential(potential: Float) -> Float {
+    sigmoid(potential)
+}
+
+fn vectorized_sigmoid_activation_of_potential(potentials: &[Float], out: &mut [Float]) {
+    for (potential, y) in potentials.iter().zip(out) {
+        *y = activation_of_sigmoid_potential(*potential)
+    }
+}
+
+fn derivative_of_sigmoid_activation_of_potential(potential: Float) -> Float {
+    let s = sigmoid(potential);
+    s * (1.0 - s)
+}
+
+fn vectorized_gradient_of_sigmoid_activation_of_potential_mut(
+    potentials: &[Float],
+    derivatives_state: &mut [Float],
+    output_gradient: &[Float],
+    input_gradient: &mut [Float],
+) {
+    use crate::linear_algebra::DiagonalMatrix;
+    for (potential, state) in potentials.iter().zip(derivatives_state.iter_mut()) {
+        *state = derivative_of_sigmoid_activation_of_potential(*potential)
+    }
+    DiagonalMatrix::new(derivatives_state).apply_mut(output_gradient, input_gradient);
+}
+
+// =====cross-entropy=====
+// takes in two probability distributions
+fn cross_entropy(p: &[Float], q: &[Float]) -> Float {
+    let mut result = 0.0;
+    for (x, y) in p.iter().zip(q) {
+        result += y * x.ln()
+    }
+    -result
+}
+
+// Second probability distribution is deterministic,
+// i.e. it deterministically outputs the same value.
+fn cross_entropy_simple(p: &[Float], q: usize) -> Float {
+    -p[q].ln()
+}
+
+fn cross_entropy_derivative_mut(p: &[Float], q: &[Float], out: &mut [Float]) {
+    for ((a, b), c) in p.iter().zip(q).zip(out) {
+        *c = -b / a
+    }
+}
+
+// Second probability distribution is deterministic,
+// i.e. it deterministically outputs the same value.
+pub fn cross_entropy_derivative_simple(p: &[Float], q: usize, out: &mut [Float]) {
+    // TODO: Do we really need to reset everything to besides the q-th index?
+    for a in out.iter_mut() {
+        *a = 0.0;
+    }
+    out[q] = -1.0 / p[q];
+}
+
+fn softmax_mut(input: &[Float], out: &mut [Float]) {
+    let mut s = 0.0;
+    for (x, y) in input.iter().zip(out.iter_mut()) {
+        let e = x.exp();
+        *y = e;
+        s += e
+    }
+
+    for y in out {
+        *y /= s
+    }
+}
+
+fn softmax_gradient_mut(
+    softmax_output: &[Float],
+    gradient_output: &[Float],
+    gradient_input: &mut [Float],
+) {
+    for (j, dx) in gradient_input.iter_mut().enumerate() {
+        *dx = 0.0;
+        for (i, dy) in gradient_output.iter().enumerate() {
+            // Note that the gradient matrix is symmetric, so don't worry about the order of
+            // indices
+            *dx += softmax_output[i] * (if i == j { 1.0 } else { 0.0 } - softmax_output[j]) * dy
+        }
+    }
+}
+
+// relu
+fn relu_mut(input: &[Float], out: &mut [Float]) {
+    for (x, y) in input.iter().zip(out.iter_mut()) {
+        *y = x.max(0.0)
+    }
+}
+
+fn relu_gradient_mut(input: &[Float], gradient_output: &[Float], gradient_input: &mut [Float]) {
+    for ((dy, dx), x) in gradient_output.iter().zip(gradient_input).zip(input) {
+        *dx = if *x > 0.0 { *dy } else { 0.0 }
+    }
+}
+
+// =====sigmoid=====
+#[derive(Debug)]
+pub struct SigmoidTransform {
+    pub weight: ColumnEfficientMatrix,
+    potential_vector: Vector,
+    derivatives_state: Vector,
+    potential_gradient: Vector,
+
+    pub weight_gradient: ColumnEfficientMatrix,
+}
+
+impl SigmoidTransform {
+    pub fn new(input_dimension: usize, output_dimension: usize) -> Self {
+        let mean = 0.0;
+        let std_dev = 1.0 / (input_dimension as Float).sqrt();
+        let normal_distr = Normal::new(mean, std_dev).unwrap();
+
+        Self {
+            weight: ColumnEfficientMatrix::random_with_normal_distribution(
+                input_dimension,
+                output_dimension,
+                normal_distr,
+            ),
+            potential_vector: Vector::zero(output_dimension),
+            derivatives_state: Vector::zero(output_dimension), // TODO: Can I get rid of this?
+            potential_gradient: Vector::zero(output_dimension),
+
+            weight_gradient: ColumnEfficientMatrix::zero(input_dimension, output_dimension),
+        }
+    }
+
+    pub fn output_mut(&mut self, input: &[Float], output: &mut [Float]) {
+        self.weight.apply_mut(input, &mut self.potential_vector[..]); // potential = W[input]
+        vectorized_sigmoid_activation_of_potential(&self.potential_vector[..], output);
+        // y = f(potential)
+    }
+
+    // Note below that (1) and (2) are independent, but they both depend on (0).
+
+    // (0)
+    pub fn potential_gradient_mut(&mut self, output_gradient: &[Float]) {
+        // updates the potential gradient
+        vectorized_gradient_of_sigmoid_activation_of_potential_mut(
+            &self.potential_vector[..],
+            &mut self.derivatives_state[..],
+            output_gradient,
+            &mut self.potential_gradient[..],
+        ); // potential_gradient = grad[f](potential)[output_gradient]
+    }
+
+    // Note that it makes sense to have the two gradients split, since for the input layer we will
+    // not need to compute the input gradient, only the weight gradient is important.
+    // WARNING: You need to call `potential_gradient_mut` before using the below function
+    // (1)
+    pub fn gradient_with_respect_to_input_mut(&self, input_gradient: &mut [Float]) {
+        // updates the input gradient
+        //
+        // Note that the first column of `self.weight` is the bias,
+        // and the previous layer doesn't care about its gradient.
+        // So we just return the gradient below the first component of the input
+        // by dropping the bias column.
+        self.weight
+            .drop_first_column_coapply_mut(&self.potential_gradient[..], input_gradient);
+        // transpose[T without the first column][potential_gradient]
+    }
+
+    // Note that the proof ensures that the `potential_gradient` has been updated
+    // WARNING: You need to call `potential_gradient_mut` before using the below function
+    // (2)
+    pub fn add_gradient_with_respect_to_weights_mut(&mut self, input: &[Float]) {
+        use crate::linear_algebra::VectorTensorCovectorMatrix;
+
+        let matrix = VectorTensorCovectorMatrix::new(&self.potential_gradient[..], input); // grad[f](potential)[output_grad] **tensor** input
+        matrix.add_to_mut(&mut self.weight_gradient);
+    }
+}
+
+// =====softmax=====
+#[derive(Debug)]
+pub struct SoftmaxTransform {
+    pub weight: ColumnEfficientMatrix,
+    potential_vector: Vector,
+    softmax_output: Vector, // Used for computation of the softmax gradient
+    potential_gradient: Vector,
+
+    pub weight_gradient: ColumnEfficientMatrix,
+}
+
+impl SoftmaxTransform {
+    pub fn new(input_dimension: usize, output_dimension: usize) -> Self {
+        let mean = 0.0;
+        let std_dev = 1.0 / (input_dimension as Float).sqrt();
+        let normal_distr = Normal::new(mean, std_dev).unwrap();
+
+        Self {
+            weight: ColumnEfficientMatrix::random_with_normal_distribution(
+                input_dimension,
+                output_dimension,
+                normal_distr,
+            ),
+            potential_vector: Vector::zero(output_dimension),
+            softmax_output: Vector::zero(output_dimension),
+            potential_gradient: Vector::zero(output_dimension),
+
+            weight_gradient: ColumnEfficientMatrix::zero(input_dimension, output_dimension),
+        }
+    }
+
+    pub fn output_mut(&mut self, input: &[Float], output: &mut [Float]) {
+        self.weight.apply_mut(input, &mut self.potential_vector[..]); // potential = W[input]
+        softmax_mut(&self.potential_vector[..], output); // y = f(potential)
+        self.softmax_output.copy_from_slice(output)
+    }
+
+    pub fn potential_gradient_mut(&mut self, output_gradient: &[Float]) {
+        softmax_gradient_mut(
+            &self.softmax_output[..],
+            output_gradient,
+            &mut self.potential_gradient[..],
+        ); // potential_gradient = grad[softmax](potential)[output_gradient]
+    }
+
+    pub fn gradient_with_respect_to_input_mut(&self, input_gradient: &mut [Float]) {
+        self.weight
+            .drop_first_column_coapply_mut(&self.potential_gradient[..], input_gradient);
+        // transpose[T without the first column][potential_gradient]
+    }
+
+    pub fn add_gradient_with_respect_to_weights_mut(&mut self, input: &[Float]) {
+        use crate::linear_algebra::VectorTensorCovectorMatrix;
+
+        let matrix = VectorTensorCovectorMatrix::new(&self.potential_gradient[..], input); // grad[f](potential)[output_grad] **tensor** input
+        matrix.add_to_mut(&mut self.weight_gradient);
+    }
+}
+
+#[derive(Debug)]
+pub struct ReluTransform {
+    pub weight: ColumnEfficientMatrix,
+    potential_vector: Vector,
+    potential_gradient: Vector,
+
+    pub weight_gradient: ColumnEfficientMatrix,
+}
+
+impl ReluTransform {
+    pub fn new(input_dimension: usize, output_dimension: usize) -> Self {
+        let mean = 0.0;
+        let std_dev = 1.0 / (input_dimension as Float).sqrt();
+        let normal_distr = Normal::new(mean, std_dev).unwrap();
+
+        Self {
+            weight: ColumnEfficientMatrix::random_with_normal_distribution(
+                input_dimension,
+                output_dimension,
+                normal_distr,
+            ),
+            potential_vector: Vector::zero(output_dimension),
+            potential_gradient: Vector::zero(output_dimension),
+
+            weight_gradient: ColumnEfficientMatrix::zero(input_dimension, output_dimension),
+        }
+    }
+
+    pub fn output_mut(&mut self, input: &[Float], output: &mut [Float]) {
+        self.weight.apply_mut(input, &mut self.potential_vector[..]); // potential = W[input]
+        relu_mut(&self.potential_vector[..], output); // y = f(potential)
+    }
+
+    pub fn potential_gradient_mut(&mut self, output_gradient: &[Float]) {
+        relu_gradient_mut(
+            &self.potential_vector[..],
+            output_gradient,
+            &mut self.potential_gradient[..],
+        ); // potential_gradient = grad[softmax](potential)[output_gradient]
+    }
+
+    pub fn gradient_with_respect_to_input_mut(&self, input_gradient: &mut [Float]) {
+        self.weight
+            .drop_first_column_coapply_mut(&self.potential_gradient[..], input_gradient);
+        // transpose[T without the first column][potential_gradient]
+    }
+
+    pub fn add_gradient_with_respect_to_weights_mut(&mut self, input: &[Float]) {
+        use crate::linear_algebra::VectorTensorCovectorMatrix;
+
+        let matrix = VectorTensorCovectorMatrix::new(&self.potential_gradient[..], input); // grad[f](potential)[output_grad] **tensor** input
+        matrix.add_to_mut(&mut self.weight_gradient);
+    }
+}