Custom TensorFlow Loss Functions for Advanced Machine Learning | by Haihan Lan | Towards Data Science

In this article, we’ll look at:
The use of custom loss functions in advanced ML applications
Defining a custom loss function and integrating to a basic Tensorflow neural net model
A brief example of knowledge distillation learning using a Gaussian Process reference applied to a few-shot learning problem
Links to my other articles:
Deep Kernel Transfer and Gaussian ProcessesRandom forestsSoftmax classificationClimate analysis
Deep Kernel Transfer and Gaussian Processes
Random forests
Softmax classification
Climate analysis
[NOTE: Tensorflow 2.0 is now released and is quite different from the 1.x version of TF that this article is based on. I am still researching TF 2.0 w.r.t custom losses and tensor manipulation as outlined in this article and may write a new article or update this one. Please keep this in mind going forward.]
Introduction
The beaten path of machine learning involves journeying to familiar landmarks and scenic locales. One set of familiar landmarks are predefined loss functions that give you a suitable loss value for the problem you are trying to optimize over. We’re familiar with the cross-entropy loss for classification and the mean squared error (MSE) or root-mean square error (RMSE) for regression problems. Popular ML packages including front-ends such as Keras and back-ends such as Tensorflow, include a set of basic loss functions for most classification and regression tasks. But off the beaten path there exist custom loss functions you may need to solve a certain problem, which are constrained only by valid tensor operations.
In Keras you can technically create your own loss function however the form of the loss function is limited to some_loss(y_true, y_pred) and only that. If you tried to add additional parameters to the loss in the form of some_loss_1(y_true, y_pred, **kwargs), Keras will throw a runtime exception and you lose the compute time that went into aggregating the datasets. There are hacks to work around this, but in general we want a scalable way to write a loss function accepting any valid arguments we pass to it, and operates on our tensors in an standard and expected way. We’ll see how to use Tensorflow directly to write a neural network from scratch and build a custom loss function to train it.
Tensorflow
Tensorflow (TF) is a symbolic and numeric computation engine that allows us to string tensors* together into computational graphs and do backpropogation over them. Keras is an API or front-end running on top of Tensorflow that conveniently packages standard constructs built using Tensorflow (such as various pre-defined neural net layers) and abstracts many of the low level mechanics of TF from the programmer (Keras can run on top of Theano as well, the same concepts apply). However, in the process of making these constructs ‘off-the-shelf’, granular level control and the power to do very specific things is lost.
*For simplicity’s sake, tensors are multi-dimensional arrays with a shape tuple like (feature_dim, n_features)
One example is the ability to define custom loss functions accepting an arbitrary number of parameters, and can compute losses with arbitrary tensors internal to the network and input tensors external to the network. Strictly speaking, a loss function in TF does not even need to be a python function, but only a valid combination of operations on TF tensor objects. The previous point is important because the power of custom losses comes from the ability to compute losses over arbitrary tensors, not strictly just your supervised target tensor and the network output tensor, in the form of (y_true, y_pred).
Before we get to custom losses, let’s briefly review a basic 2-layer dense net (MLP) and see how it’s defined and trained in TF. Although there are predefined TF layers, let’s define the layers from scratch from weights and bias tensors. We want to get familiar with the placeholder and variable tensors in TF.
That was pretty cool. The code above can be modified for multi-class classification by replacing the loss with softmax_cross_entropy_with_logits and the final sigmoid activation replaced with a tf.nn.softmax.
Next to demonstrate how to use custom loss functions with arbitrary tensors, let’s implement a knowledge distillation model which optimizes over the binary classification loss as well as a loss between the model being trained and a reference model. Knowledge distillation is a form of transfer learning where we learn with our target model (model we want to train) but also indirectly transfer knowledge representations from a reference model. We’ll use a Gaussian Process Classifier (GPC) from sklearn as our reference model. We’ll also make the problem more interesting by reducing our training data to only 1% of the samples from the 569 in the sklearn breast cancer dataset... for both the reference and the target, and train them both from scratch.* This is known as a few-shot learning problem.
*In traditional transfer learning the reference model is usually a wide and/or deep network pre-trained on many examples/classes and the target is a narrower/shallower network to be trained on few available specific examples/classes.
The loss for this knowledge distillation scheme looks something like
The binary cross entropy loss is just the regular binary classification loss, the second term involves another loss D between the outputs of the target f(x) and the reference g(x). We’ll let D be the Kullback-Leibler Divergence (DKL) between f(x) and g(x):
The DKL in a nutshell quantifies how different a distribution f is from g, in terms of information (roughly information is inversely proportional to certainty); it can be thought of as a cross entropy between distributions, and is an asymmetric loss that can take negative values*. By minimizing the DKL between f and g we basically want to to increase the information content of f relative to g. When f and g have the same amount of information, the log term above is 0 and the DKL loss is 0 as well. Using the DKL as the loss makes sense when using GPCs as a reference model, since when making predictions form the GPC, we are sampling from its posterior distribution (a softmax) and although our neural net is a crude approximation of this posterior, it is also a distribution.
*In our implementation below, we take abs( D_KL( f(x), g(x) ) ) in our combined loss. Theoretically D_KL will always be non-negative due to the log-sum inequality but computing D_KL with round-off errors on a real computer can lead to negative values.
Notice how we need to now feed external inputs g(x) to our loss. In Keras this process is contrived and unscalable. But in TF, it is as simple as creating a new placeholder tensor, adding the necessary terms to the combined loss, and feeding the inputs when running a training or prediction session.
Having only 5 training samples, the above example converges faster and gives better raw test accuracy than a net where the DKL loss is set to zero (i.e. no transfer learning). Note, test set imbalance was not accounted for! The astute reader should add an F1 score to the accuracy readout. Thankfully our pseudo-random sample of the dataset gives a 2:3 split of positive vs negative classes.
Also interesting to note is the softmax outputs of the transfer learning model tested on 100 examples from the holdout:
[0.5134167 ] [0.5767678 ] [0.5299671 ] [0.529941  ] [0.51807505] [0.4615707 ] [0.61761355] [0.5744497 ] [0.6092696 ] [0.55092424] [0.5866923 ] [0.5522269 ] [0.5679551 ]
Versus the non-transfer learning model
[0.44836044] [0.99457294] [0.47165167] [0.573235  ] [0.94637066] [0.00909297] [0.99778605] [0.99487936] [0.99742365] [0.96588767] [0.99646676] [0.9843067 ] [0.99225134]
The predictions of the transfer learning model reflect the uncertainty given the limited information it was trained on. There is a very good reason why a GPC was chosen as the reference model, after all how can we be so sure a new patient has cancer when we trained a classifier on only 5 examples from scratch?
Conclusion
We saw how to implement a neural net from scratch in TF, how to combine tensor operations into loss functions, and touched on an interesting application of transfer learning. In general TF is far more flexible for the Data Scientist in advanced or niche supervised learning applications. If you enjoyed reading this article and using the code, please check out my other articles!