TensorFlow Integration
In reservr: Fit Distributions and Neural Networks to Censored and Truncated Data

knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>"
)

Overview

While {reservr} is capable of fitting distributions to censored and truncated observations, it does not directly allow modelling the influence of exogenous variables observed alongside the primary outcome. This is where the integration with TensorFlow comes in.

The TensorFlow integration allows to fit a neural network simultaneously to all parameters of a distribution while taking exogenous variables into account.

{reservr} accepts all partial tensorflow networks which yield a single arbitrary-dimension rank 2 tensor (e.g. any dense layer) as output and can connect suitable layers to this intermediate output such that the complete network predicts the parameters of any pre-specified distribution family.

It also dynamically compiles a suitable conditional likelihood based loss, depending on the type of problem (censoring, truncation), which can be optimized using the keras3::fit implementation out-of-the box. This means there is full flexibility with respect to callbacks, optimizers, mini-batching, etc.

library(reservr)
library(tensorflow)
library(keras3)
library(tibble)
library(ggplot2)

A simple linear model

The following example will show the code necessary to fit a simple model with the same assumptions as OLS to data. As a true relationship we use $y = 2 x + \epsilon$ with $\epsilon \sim \mathcal{N}(0, 1)$. We will not use censoring or truncation.

if (reticulate::py_module_available("keras")) {
  set.seed(1431L)
  tensorflow::set_random_seed(1432L)

  dataset <- tibble(
    x = runif(100, min = 10, max = 20),
    y = 2 * x + rnorm(100)
  )

  print(qplot(x, y, data = dataset))

  # Specify distributional assumption of OLS:
  dist <- dist_normal(sd = 1.0) # OLS assumption: homoskedasticity

  # Optional: Compute a global fit
  global_fit <- fit(dist, dataset$y)

  # Define a neural network
  nnet_input <- layer_input(shape = 1L, name = "x_input")
  # in practice, this would be deeper
  nnet_output <- nnet_input

  optimizer <- optimizer_adam(learning_rate = 0.1)

  nnet <- tf_compile_model(
    inputs = list(nnet_input),
    intermediate_output = nnet_output,
    dist = dist,
    optimizer = optimizer,
    censoring = FALSE, # Turn off unnecessary features for this problem
    truncation = FALSE
  )

  nnet_fit <- fit(nnet, x = dataset$x, y = dataset$y, epochs = 100L, batch_size = 100L, shuffle = FALSE)

  # Fix weird behavior of keras3
  nnet_fit$params$epochs <- max(nnet_fit$params$epochs, length(nnet_fit$metrics$loss))
  print(plot(nnet_fit))

  pred_params <- predict(nnet, data = list(as_tensor(dataset$x, config_floatx())))

  lm_fit <- lm(y ~ x, data = dataset)

  dataset$y_pred <- pred_params$mean
  dataset$y_lm <- predict(lm_fit, newdata = dataset, type = "response")

  p <- ggplot(dataset, aes(x = x, y = y)) %+%
    geom_point() %+%
    geom_line(aes(y = y_pred)) %+%
    geom_line(aes(y = y_lm), linetype = 2L)
  print(p)

  coef_nnet <- rev(as.numeric(nnet$model$get_weights()))
  coef_lm <- coef(lm_fit)

  print(coef_nnet)
  print(coef_lm)
}