That OpenAI work is more elaborate. It trains an additional network in such a way that it encodes what GPT is doing in terms of activations, but in a more interpretable way (hopefully). Here, as far as I can tell, it's visualizing the activation of the attention layers directly.