What is the correct code for AllenNLP/NER task?

[bratao]

Hello @amirgholami ,

I´m super excited about this optimizer. Thank you!

I want to use it in a NER task using AllenNLP. But I´m confused because the code differs between the image_classification and transformer examples.

At https://github.com/amirgholami/adahessian/blob/5c176cdcbeacff1d9edfc77062d0bc7594f326a9/image_classification/optim_adahessian.py in function get_trace, we have:

hutchinson_trace = []
        for hv, vi in zip(hvs, v):
            param_size = hv.size()
            if len(param_size) <= 2:  # for 0/1/2D tensor
                tmp_output = torch.abs(hv * vi)
                hutchinson_trace.append(tmp_output) # Hessian diagonal block size is 1 here.
            elif len(param_size) == 4:  # Conv kernel
                tmp_output = torch.abs(torch.sum(torch.abs(
                    hv * vi), dim=[2, 3], keepdim=True)) / vi[0, 1].numel() # Hessian diagonal block size is 9 here: torch.sum() reduces the dim 2/3.
                hutchinson_trace.append(tmp_output)

While in https://github.com/amirgholami/adahessian/blob/bd9f5a6760bf1ba4474e2e8a5fad237a1577d989/transformer/fairseq/optim/adahessian.py we have:

hutchinson_trace = []
        for hv, vi in zip(hvs, v):
            param_size = hv.size()
            if len(param_size) <= 1: # for Bias and LN 
                tmp_output = torch.abs( hv * vi)  + 0.
                hutchinson_trace.append( tmp_output )
            elif len(param_size) == 2: # Matrix
                tmp_output1 = torch.abs((hv * vi + 0.)).view(-1, self.block_length) # faltten to the N times self.block_length
                tmp_output2 = torch.abs(torch.sum(tmp_output1, dim=[1])).view(-1) / float(self.block_length)
                tmp_output3 = tmp_output2.repeat_interleave(self.block_length).view(param_size)
                hutchinson_trace.append(tmp_output3)

Which one should I choose?

In my NLP task I have parameters with sizes varying between 1 and 4. For parameters with size 3 , neither would match it in the loop. Is this correct?

[amirgholami]

Hi @bratao,

Thanks for your interest. Both of these functions do the same thing which is to compute the Hessian diagonal. The only difference is the shape of the model parameters. The first code is written for convolutional neural networks that have four dimension (C_inC_outKK), or 2 dimensional as in the case of a FC layer.

The second code is specifically for transformers and in particular for Bias and LN which have a dimension size of 1, and the attention layers.

Could you please let me know for which layer type you have a parameter size of 3?

[bratao]

It is a CNN for getting a char representation of a token

params = Params(
            {
                "embedding": {"embedding_dim": 16, "vocab_namespace": "token_characters"},
                "encoder": {
                    "type": "cnn",
                    "embedding_dim": 16,
                    "num_filters": 128,
                    "ngram_filter_sizes": [3],
                    "conv_layer_activation": "relu",
                },
            }
        )

I mixed the two versions in the function abobe:

    def get_trace(self, gradsH):
        """
        compute the Hessian vector product with a random vector v, at the current gradient point,
        i.e., compute the gradient of <gradsH,v>.
        :param gradsH: a list of torch variables
        :return: a list of torch tensors
        """

        params = self.param_groups[0]["params"]
        params = list(filter(lambda x: x.requires_grad, params))

        v = [torch.randint_like(p, high=2, device=self.device) for p in params]
        for v_i in v:
            v_i[v_i < 0.5] = -1
            v_i[v_i >= 0.5] = 1

        hvs = torch.autograd.grad(
            gradsH, params, grad_outputs=v, only_inputs=True, retain_graph=True
        )

        hutchinson_trace = []
        for hv, vi in zip(hvs, v):
            param_size = hv.size()
            if len(param_size) <= 2:  # for Bias and LN
                tmp_output = torch.abs(hv * vi) + 0.0
                hutchinson_trace.append(tmp_output)
            else:  # Matrix
                tmp_output1 = torch.abs((hv * vi + 0.0)).view(
                    -1, self.block_length
                )  # flatten to the N times self.block_length
                tmp_output2 = torch.abs(torch.sum(tmp_output1, dim=[1])).view(-1) / float(
                    self.block_length
                )
                tmp_output3 = tmp_output2.repeat_interleave(self.block_length).view(param_size)
                hutchinson_trace.append(tmp_output3)

        return hutchinson_trace

Apparently it works. All my test suite pass. Despite being slower to converge than the Ranger optimizer, it do not requires an adjustment of the LR, that is an great trade-off.

[amirgholami]

I am happy to hear that you are finding good use for it. We are actually observing strong improvement for NLP tasks, where the average GLUE score significantly increases with AdaHESSIAN. We will soon update the paper with these results. We would also like to hear more about the details of your use case if you would like to share them.

Regarding the implementation, we perform the block averaging on the convolution dimensions. Please see: https://github.com/amirgholami/adahessian/blob/5c176cdcbeacff1d9edfc77062d0bc7594f326a9/image_classification/optim_adahessian.py#L92

Specifically note that the averaging for convolution filters is occurring across dim=[2,3] which is the filter size dimension. You may get better performance by doing that here as well. For example, for a 3x3 conv filter, the block averaging happens across groups of 9 convolution parameters.

Also regarding speed you may want to try the delayed hutchinson step calculation so you are computing Hessian diagonal every other iteration. But even though AdaHessian is a little slower, it gives a good tradeoff with reduced hyperparameter tuning.

[bratao]

I will try @amirgholami , But I just got an error when I moved to my production cluster: RuntimeError: derivative for _cudnn_rnn_backward is not implemented

Apparently I will need to use without cudnn. Is that right? 😢

[amirgholami]

No it does work with cudnn. Is there a sample code that we can take a look at? I would like to see exactly how the convolution is being applied in your code. Based on the above snippet the block averaging should happen across ngram_filter_sizes

[bratao]

@amirgholami Here is a repo I did with a regular BI-LSTM-CRF model on the CONLL-2003 task.

https://github.com/bratao/ner_adahessian

It compares with the Ranger optimizer

[yaozhewei]

Hi Bratao,

We added instructions to support different types of kernels. Please let us know if this help solve your problem.

BTW: Currently, PyTorch does not support second-order derivative for RNN type of layers (like LSTM, GRU, RNN). Therefore, if you are asking how to use AdaHessian for those models, there is no solution yet.

Best,