Author: Pablo Morala
Date: 2024-01-17

# Supported DL frameworks

#### Pablo Morala

#### 2024-01-17

In this vignette we will explore the supported deep learning frameworks in `nn2poly` and how can they be used. Currently supported frameworks are `keras`/`tensorflow` and `luz`/`torch`. The benefit of using them is that we provide an easy to use implementation of the weight constraints needed during training to avoid problems in the Taylor expansion carried out by `nn2poly`. Furthermore, models created with these frameworks are directly supported by the `nn2poly()` function, allowing then to skip the need to manually create an object with weights and activation functions as explained in `vignette("nn2poly-01-introduction")`.

If you are interested in having support for another deep learning framework, please open an issue in our GitHub repo!

> Note: Using the provided constraints when training a model increases the number of needed epochs to converge, which is an expected trade off between learning speed and control over the weights for easier explainability.

# Data generation

As our goal here is to show how to impose the needed constraints for each framework, we will replicate the polynomial data generation from `vignette("nn2poly-01-introduction")` and solve later the same regression problem with the different deep learning frameworks. Refer to that vignette for more details.

    library(nn2poly)
    set.seed(42)

Create the polynomial \(4x_1 - 3 x_2x_3\):

    polynomial <- list()
    polynomial$labels <- list(c(1), c(2,3))
    polynomial$values <- c(4,-3)

Generate the data:

    # Define number of variables and sample size
    p <- 3
    n_sample <- 500

    # Predictor variables
    X <- matrix(0,n_sample,p)
    for (i in 1:p){
      X[,i] <- rnorm(n = n_sample,0,1)
    }

    # Response variable + small error term
    Y <- nn2poly:::eval_poly(poly = polynomial, newdata = X) +
      stats::rnorm(n_sample, 0, 0.1)

    # Store all as a data frame
    data <- as.data.frame(cbind(X, Y))
    head(data)
    #>           V1           V2         V3          Y
    #> 1  1.3709584  1.029140719  2.3250585 -1.7547416
    #> 2 -0.5646982  0.914774868  0.5241222 -3.7107357
    #> 3  0.3631284 -0.002456267  0.9707334  1.3609395
    #> 4  0.6328626  0.136009552  0.3769734  2.4608270
    #> 5  0.4042683 -0.720153545 -0.9959334 -0.6141076
    #> 6 -0.1061245 -0.198124330 -0.5974829 -0.7455793

Scale the data to have everything in the \([-1,1]\) interval and divide it in train and test.

    # Data scaling to [-1,1]
    maxs <- apply(data, 2, max)
    mins <- apply(data, 2, min)
    data <- as.data.frame(scale(data, center = mins + (maxs - mins) / 2, scale = (maxs - mins) / 2))

    # Divide in train (0.75) and test (0.25)
    index <- sample(1:nrow(data), round(0.75 * nrow(data)))
    train <- data[index, ]
    test <- data[-index, ]

    train_x <- as.matrix(train[,-(p+1)])
    train_y <- as.matrix(train[,(p+1)])

    test_x <- as.matrix(test[,-(p+1)])
    test_y <- as.matrix(test[,(p+1)])

# Supported frameworks examples

With our common data generated, we are now ready to delve into the different deep learning frameworks. In each case we will solve the same regression problem with the previously generated polynomial data, with two neural networks for each framework. These neural networks will have the same structure but one will be constrained (`nn_con`) and the other one unconstrained (`nn_uncon`) to showcase their differences.

## `keras`/`tensorflow`

In this section we will show how to use `nn2poly` with a `keras` neural network and how to impose the needed weight constraints during training.

    library(nn2poly)
    library(keras)

    # This sets all needed seeds
    tensorflow::set_random_seed(42)

### Model definition

First of all we will set the structure of our neural networks using a sequential model in `keras`. This following function will be used to create both the constrained and unconstrained neural networks in the `keras` example.

    keras_model <- function() {
      tensorflow::set_random_seed(42)

      nn <- keras::keras_model_sequential()
      nn <- keras::layer_dense(nn, units = 100, activation = "tanh", input_shape = p)
      nn <- keras::layer_dense(nn, units = 100, activation = "tanh")
      nn <- keras::layer_dense(nn, units = 100, activation = "tanh")
      nn <- keras::layer_dense(nn, units = 1, activation = "linear")

      nn
    }

Now we define both NNs. The only needed step to impose our implemented weight constraints on a `keras` model is to use the function `add_constraints()` and then train it as usually. The constraints currently accept two possible types, `"l1_norm"` and `"l2_norm"`. In both cases, the norm for each weight vector (including the bias) incident on a neuron will be constrained to be less or equal than 1. Using data scaled to the \([-1,1]\) interval, the l1-norm is the one that guarantees best results.

Also, it is important to note that the weight constraints are imposed on all layers except for the last one, which is expected to be linear and not need a Taylor expansion.

> Note: Our implementation differs slightly from the constraints provided by `keras` because in our implementation we join the bias with the rest of the weights while `keras` default constraints allow only for bias or kernel constraints. Our implementation uses a custom callback that is applied at the end of each batch.

    nn_uncon <- keras_model()

    nn_con <- keras_model()
    nn_con <- add_constraints(nn_con, constraint_type = "l1_norm")

### Training

Now we can compile and train both NNs using standard the `keras` approach. Note that we have to increase the number of needed epochs with the constrained NN to learn properly.

    compile(nn_uncon,
            loss = "mse",
            optimizer = optimizer_adam(),
            metrics = "mse")

    history1 <- fit(nn_uncon,
                   train_x,
                   train_y,
                   verbose = 0,
                   epochs = 150,
                   batch_size = 50,
                   validation_split = 0.2
    )

    plot(history1)

    compile(nn_con,
            loss = "mse",
            optimizer = optimizer_adam(),
            metrics = "mse")

    history2 <- fit(nn_con,
                   train_x,
                   train_y,
                   verbose = 0,
                   epochs = 700,
                   batch_size = 50,
                   validation_split = 0.2
    )

    plot(history2)

### NN predictions

We can visualize the NN predictions vs the original Y values for both neural networks and observe how both of them provide accurate predictions (the values fall near the “perfect” diagonal red line).

    # Obtain the predicted values with the NN to compare them
    prediction_nn_uncon <- predict(nn_uncon, test_x)
    #> 4/4 - 0s - 48ms/epoch - 12ms/step

    # Diagonal plot implemented in the package to quickly visualize and compare predictions
    nn2poly:::plot_diagonal(x_axis =  prediction_nn_uncon, y_axis =  test_y, xlab = "Unconstrained NN prediction", ylab = "Original Y")

    # Obtain the predicted values with the NN to compare them
    prediction_nn_con <- predict(nn_con, test_x)
    #> 4/4 - 0s - 91ms/epoch - 23ms/step

    # Diagonal plot implemented in the package to quickly visualize and compare predictions
    nn2poly:::plot_diagonal(x_axis =  prediction_nn_con, y_axis =  test_y, xlab = "Constrained NN prediction", ylab = "Original Y")

### Using nn2poly

After the NNs have been trained, we can directly call `nn2poly` on the `keras` model.

    # Polynomial for nn_uncon
    final_poly_uncon <- nn2poly(object = nn_uncon,
                          max_order = 3)

    # Polynomial for nn_con
    final_poly_con <- nn2poly(object = nn_con,
                          max_order = 3)

We can visualize the polynomial predictions versus NN predictions

    # Obtain the predicted values for the test data with our two polynomials
    prediction_poly_uncon <- predict(object = final_poly_uncon, newdata = test_x)
    prediction_poly_con <- predict(object = final_poly_con, newdata = test_x)

    nn2poly:::plot_diagonal(x_axis =  prediction_nn_uncon, y_axis =  prediction_poly_uncon, xlab = "NN prediction", ylab = "Polynomial prediction") + ggplot2::ggtitle("Polynomial for nn_uncon")

    nn2poly:::plot_diagonal(x_axis =  prediction_nn_con, y_axis =  prediction_poly_con, xlab = "NN prediction", ylab = "Polynomial prediction") + ggplot2::ggtitle("Polynomial for nn_con")

We can clearly see how the constrained NN allows to obtain a close approximation with `nn2poly` while the polynomial obtained for the unconstrained one is not a good representation.

We can also represent the obtained polynomial coefficients and observe how the constrained case clearly has terms `2,3` and `1` with the same sign as the original polynomial while the rest are close to 0.

> Note: The coefficients values are not in the same scale as the original polynomial due to the fact that we have scaled all the data before training, even the response variable Y. Furthermore, as data has been scaled to the \([-1,1]\) interval, interactions of order 2 or higher would usually need a higher absolute value than the lower order coefficients to be more relevant

    plot(final_poly_uncon, n = 8)

    plot(final_poly_con, n = 8)

## `luz`/`torch`

In this section we will show how to use `nn2poly` with a `torch` neural network, built with its higher level API `luz`, and how to impose the needed weight constraints during training. Furthermore, we will explain how to use `luz_model_sequential()`, a helper needed to create a `torch` model in an adequate manner so that it can be easily recognized by `nn2poly`.

    library(torch)
    library(luz)

    set.seed(42)

### Data loader for `torch`

In this framework, we need to manipulate a bit our data to use it as input for a `torch` model. This can be done dividing our `train_x` data intro only train and validation matrices and using `luz::as_dataloader()`

    # Divide in only train and validation
    all_indices   <- 1:nrow(train_x)
    only_train_indices <- sample(all_indices, size = round(nrow(train_x)) * 0.8)
    val_indices   <- setdiff(all_indices, only_train_indices)

    # Create lists with x and y values to feed luz::as_dataloader()
    only_train_x <- as.matrix(train_x[only_train_indices,])
    only_train_y <- as.matrix(train_y[only_train_indices,])
    val_x <- as.matrix(train_x[val_indices,])
    val_y <- as.matrix(train_y[val_indices,])

    only_train_list <- list(x = only_train_x, y = only_train_y)
    val_list <- list(x = val_x, y = val_y)

    torch_data <- list(
      train = luz::as_dataloader(only_train_list, batch_size = 50, shuffle = TRUE),
      valid = luz::as_dataloader(val_list, batch_size = 50)
    )

### Model definition

Now we will set the structure of our neural networks using a sequential model in `torch`. To do so we have implemented the helper `luz_model_sequential()`, which lets the user stack linear layers in a similar way as `keras` and allows `nn2poly()` to be directly used in the model without the need to extract weights and activation functions manually as in the default example of `vignette("nn2poly-01-introduction")`.

This following function will be used to create both the constrained and unconstrained neural networks in the `torch` example.

    luz_nn <- function() {
      torch::torch_manual_seed(42)

      luz_model_sequential(
        torch::nn_linear(p,100),
        torch::nn_tanh(),
        torch::nn_linear(100,100),
        torch::nn_tanh(),
        torch::nn_linear(100,100),
        torch::nn_tanh(),
        torch::nn_linear(100,1)
      )
    }

Now we define both NNs. In this case, differing from the `keras` example, we will impose the constraints later to follow the same approach as in `luz` documentation. However, the use of `add_constraints()` is analogous.

    nn_uncon <- luz_nn()
    nn_con <- luz_nn()

### Training

First we train the unconstrained NN using ans standard `luz` approach.

    fitted_uncon <- nn_uncon %>%
        luz::setup(
          loss = torch::nn_mse_loss(),
          optimizer = torch::optim_adam,
          metrics = list(
            luz::luz_metric_mse()
          )
        ) %>%
        luz::fit(torch_data$train, epochs = 50, valid_data = torch_data$valid)

    fitted_uncon %>% plot()

Then, we train the constrained NN by using the function function `add_constraints()` before calling `fit()`. The constraints currently accept two possible types, `"l1_norm"` and `"l2_norm"`. In both cases, the norm for each weight vector (including the bias) incident on a neuron will be constrained to be less or equal than 1. Using data scaled to the \([-1,1]\) interval, the l1-norm is the one that guarantees best results.

Also, it is important to note that the weight constraints are imposed on all layers except for the last one, which is expected to be linear and not need a Taylor expansion.

Note that we have to increase the number of needed epochs with the constrained NN to learn properly.

    fitted_con <- nn_con %>%
      luz::setup(
        loss = torch::nn_mse_loss(),
        optimizer = torch::optim_adam,
      ) %>%
      add_constraints("l1_norm") %>%
      fit(torch_data$train, epochs = 700, valid_data = torch_data$valid)

    fitted_con %>% plot()

> Note: Our implementation uses a custom callback that is applied at the end of each batch.

### NN predictions

We can visualize the NN predictions vs the original Y values for both neural networks and observe how both of them provide accurate predictions (the values fall near the “perfect” diagonal red line).

    # Obtain the predicted values with the NN to compare them
    prediction_NN_uncon <- as.array(predict(fitted_uncon, test_x))

    # Diagonal plot implemented in the package to quickly visualize and compare predictions
    nn2poly:::plot_diagonal(x_axis =  prediction_NN_uncon, y_axis =  test_y, xlab = "Unconstrained NN prediction", ylab = "Original Y")

    # Obtain the predicted values with the NN to compare them
    prediction_NN_con <- as.array(predict(fitted_con, test_x))

    # Diagonal plot implemented in the package to quickly visualize and compare predictions
    nn2poly:::plot_diagonal(x_axis =  prediction_NN_con, y_axis =  test_y, xlab = "Constrained NN prediction", ylab = "Original Y")

### Using nn2poly

After the NNs have been trained, we can directly call `nn2poly` on the `luz` model.

    # Polynomial for nn_uncon
    final_poly_uncon <- nn2poly(object = fitted_uncon,
                          max_order = 3)

    # Polynomial for nn_con
    final_poly_con <- nn2poly(object = fitted_con,
                          max_order = 3)

We can visualize the polynomial predictions versus NN predictions

    # Obtain the predicted values for the test data with our two polynomials
    prediction_poly_uncon <- predict(object = final_poly_uncon, newdata = test_x)
    prediction_poly_con <- predict(object = final_poly_con, newdata = test_x)

    nn2poly:::plot_diagonal(x_axis =  prediction_nn_uncon, y_axis =  prediction_poly_uncon, xlab = "NN prediction", ylab = "Polynomial prediction") + ggplot2::ggtitle("Polynomial for nn_uncon")

    nn2poly:::plot_diagonal(x_axis =  prediction_nn_con, y_axis =  prediction_poly_con, xlab = "NN prediction", ylab = "Polynomial prediction") + ggplot2::ggtitle("Polynomial for nn_con")

We can clearly see how the constrained NN allows to obtain a close approximation with `nn2poly` while the polynomial obtained for the unconstrained one is not a good representation.

We can also represent the obtained polynomial coefficients and observe how the constrained case clearly has terms `2,3` and `1` with the same sign as the original polynomial while the rest are close to 0.

> Note: The coefficients values are not in the same scale as the original polynomial due to the fact that we have scaled all the data before training, even the response variable Y. Furthermore, as data has been scaled to the \([-1,1]\) interval, interactions of order 2 or higher would usually need a higher absolute value than the lower order coefficients to be more relevant

    plot(final_poly_uncon, n = 8)

    plot(final_poly_con, n = 8)