Debugging using Orthogonal Gradient Descent

While the field of machine learning has seen incredible progress in the past decade, in our view, there is an important problem that has not received enough traction: the problem of “debugging" machine learning models. That is, when it is discovered that a trained model is partially faulty – that it responds to a portion of the inputs in an incorrect manner, perhaps due to a dataset bias or due to insufficient training data – a very common method of debugging is to simply train the model from scratch. While this naive approach works well, retraining models can be prohibitively expensive, with many of the state-of-the art models requiring several weeks and multiple GPUs to train. Finances and the time overhead aside, as a matter of principle, starting from scratch is unsatisfactory. For instance, we rarely address bugs in our mathematical models or standard computer code by throwing everything away and starting with a clean slate – and neither do we abandon a child every time they misbehave or say unpleasant things!

With the above concerns in mind, we explore the problem of debugging machine learning models without having to train the models from scratch. We do so by basing our approach on the claim that the debugging problem can be treated as a two-task continual learning problem. More specifically, we employ a modified version of a continual learning method known as Orthogonal Gradient Descent [4], and using a simple convolutional neural network trained on a (purposefully) corrupted version of the MNIST dataset we show that (1) it is possible to unlearn the way the model responds to the misclassified data points without affecting the general performance; and (2) on top of unlearning, it is also possible to relearn the right responses to misclassified inputs and attain the maximum possible test accuracy (with the simple network), without training from scratch.

In this section we talk about Orthogonal Gradient Descent (OGD), the method that underlies our approach. While continual learning, in general, focuses on the problem of sequential learning of many tasks $\{\mathcal{T}_{1},\mathcal{T}_{2},...,\mathcal{T}_{n}\}$, here we discuss the method with just two-tasks in mind, $\{\mathcal{T}_{1},\mathcal{T}_{2}\}$. We use $f(x;\theta)$ to denote the neural network, where $x\in\mathbb{R}^{d}$ is the input to the network and $\theta$ are the weights that parametrize the model. The network takes in an input $x$ and outputs its prediction $y\in\mathbb{R}^{c}$. In case of classification problems, $f_{j}(x;\theta)$ denotes the $j$-th logit associated to the $j$-th class. The loss function (for a single task) is denoted as

where the sum is over the training data and $\tilde{y}$ denotes the ground truth.

When discussing OGD, it is crucial to distinguish between two types of gradients: first, we have the gradient of the neural network, $\nabla_{\theta}f(x;\theta)$, and then we have the gradient of the loss function, $\nabla_{\theta}\mathcal{L}(\theta)$, which is the one encountered more frequently in machine learning applications as this is used to minimize the loss.

With the preliminaries established, we can now describe the OGD algorithm in some detail. Please refer to the Appendix (1) or the original paper for more information [4].

1. Computing Model Gradients After training the model on the first task in the usual manner, we first compute the model gradients of the class that corresponds to the ground truth label, that is, if an $x_{i}$ belongs to the $k$-th class, then we compute and store only $f_{k}(x_{i};\theta)$. On the MNIST dataset, the authors recommend computing the gradients at around 300 data points randomly sampled from the training set and thus the memory requirements are reasonable – note that the dataset has about 60k training points.

2. Orthogonal Basis Once we have the model gradients $S$ = {$\nabla_{\theta}f_{k}(x_{i};\theta)$}, we then compute the orthogonal basis for $S$ using the Gram-Schmidt method on all gradients. From here on, we take $S$ to denote the orthonormal set.

3. Computing Loss Gradients Third, we compute the gradients of the loss function $\mathcal{L}$ on the data points that belong to the second task $\mathcal{T}_{2}$:

4. Projection We then project the gradient of the loss function computed in the previous step onto the subspace that is orthogonal to the span of the model gradients

5. Parameter Update Finally, we update the parameters of the model using the new projected vector:

where we set $\eta$ is the step-size. In our experiments we use $\eta=0.001$.

We repeat steps 3-5 until the model converges. In this report, we make slight modifications to the original approach and the modified approach is described in later sections.

The problems discussed in this report are related to the ones explored in the field of machine unlearning [2]. This nascent field has been motivated in part by the recently introduced privacy regulations such as the General Data Protection Regulation in the European Union[11], the California Consumer Privacy Act in the United States, and Consumer Privacy Protection Act in Canada. While there are many popular approaches to data deletion [1, 6, 9, 5, 12, 3], our approach is, at a high level, closely related to the one presented in [7]; their method is based on a continual learning algorithm called Elastic Weight Consolidation [10], which uses the Fisher Information Matrix [8] to assign importance to the learned weights.

In the following section we employ the MNIST dataset to perform two simple experiments to demonstrate unlearning and relearning. Unless otherwise stated, we use the standard train/test split of 60k/10k. Throughout, we use a simple convolutional neural network architecture that has around 20k parameters, which, when trained on 60k images, attains an accuracy of 98% on the test set. We do not worry about tuning hyperparameters.

We present two experiments supporting the hypothesis that the problem of debugging can be treated as a continual learning problem. First, we use the phenomenon of catastrophic forgetting to our advantage, and show that we can in fact do targeted unlearning. The model forgets to (mis)-classify the images containing the digits ‘2’ as three and vice-versa. However, when we compare the average (highest) confidence with which the model classifies these images before and after the unlearning, we did not notice a drastic difference: 0.98 vs 0.78 – when the weights are randomly initialized the confidence is around 0.12. In the ideal case, not only would the model unlearn, but the confidence of the predictions would also look more like the randomly initialized model. Second, we present an experiment on relearning and demonstrate that the model can successfully relearn the right behaviour, reaching the maximum possible accuracy on the dataset, given the model architecture.

We believe that it would be interesting to apply these methods to deal with problems where the model has a lot more parameters and/or where the dataset is much larger. We have also considered the possibility of using modified OGD (like in the unlearning case) to increase the robustness of models against targeted adversarial attacks.