OpenAI, Week 19-20 // Empirically Comparing Two Models

When writing my own version of a workhorse model that has been implemented in a large, well-tested public codebase like PyTorch, I like to test mine via direct comparison. I've found it really helpful for giving myself peace of mind that e.g. my ResNet and Transformer implementations don't have any silly bugs in them. Here, I'll briefly discuss my process.

The first thing I like to do is test that the two models have the same number of parameters (given identical architectural configurations, of course). Here is a little script:

def print_num_params(model):
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    print(params)

print_num_params(my_model)
print_num_params(torch_model)

After that basic test is passed, I want to verify that the models are doing the same thing. I force the parameters in both models to take the same values, then pass the same random input tensor through each and take the ratio of their outputs. This ratio should be all ones, or very close. I usually print the min/mean/max or plot a histogram (you may have to flatten the output tensor first). 

To be safe, put the models in evaluation mode beforehand. This should, in theory, disable stochastic components like dropout and normalization. I'll talk about that further down.

Here is a little script to set the parameters:

def set_model_params(model, value):
    for p in model.parameters():
        p.data.fill_(value)

set_model_params(my_model, 0.5)
set_model_params(torch_model. 0.5)

Or, if you want to set the weights and biases to different values (be careful with this, of course - pay attention to how the weights and biases are labelled), you could do something like this:

def set_model_params(model, weight_value, bias_value):
    for name, parameter in model.named_parameters():
        if name.endswith('weight'):
             parameter.data.fill_(weight_value) 
        if name.endswith('bias'):
             parameter.data.fill_(bias_value) 

set_model_params(my_model, 0.5, -1.2)
set_model_params(torch_model. 0.5, -1.2)

You may be tempted to set the weights to 1 and the biases to 0, like a basic initialization scheme, but I've found this sometimes leads to false positives of equivalence. Go for random numbers. Better yet, try a few different random numbers, to make absolutely sure!

A note about stochasticitiy

If you want to be extra careful about potential stochasticity, disable dropout or set the probability to 0.0 in both models. When in doubt, remove any batch/layer normalization or replace it with a dummy version that just scales the input by a fixed number each time.

A good test to see if you've removed the stochasticity is to create two instances of the same model (yours or the one you're comparing to) and take that ratio. First of all, it should be ones, and secondly, it should not fluctuate as you repeat the experiment. For that matter, the ratio between your model and the model you're comparing to also shouldn't fluctuate - whether it's right or wrong.

All of that said, I've run into persistent issues caused by layer normalization even after doing all of these steps. It has less to do with stochasticity and more to do with numerical accuracy, but it's worth mentioning in the context of trying to get equivalence. When implementing my own Transformer, I found that I got some diverging values due to a small rounding difference that happens in the decoder stack. Because of the layer norms, these small rounding differences grow into large outliers in the output ratio. You can try to mitigate it by tuning the stability term $\epsilon$ in the layer norm, or rounding to a fixed precision in both versions of the code.

That's all for now, I hope this was helpful to others who are implementing their own models!