Applying adversarial_bias_mitigator to span-BERT

[alicesaunders]

I was hoping to apply adversarial debiasing to the task of coreference resolution. I was wondering if it is possible to do this through AllenNLP, using the spanBERT model and the adversarial_bias_mitigator model wrapper. Can these two be combined?

[epwalsh]

@ArjunSubramonian may know the answer to this

[ArjunSubramonian]

Yes, these two can be combined! You have to implement an adversary that takes as input coreference probabilities. For example, if you're using WinoBias-style (https://uclanlp.github.io/corefBias/overview) sentences, e.g. "The physician hired the secretary because she was overwhelmed with clients", the predictor, i.e. coref-spanBERT, would output coreference probabilities and the adversary would take as input the probabilities of "she" coreferring with "physician" and "secretary" and attempt to recover the pronoun. In this way, the predictor can be finetuned to produce more fair coreference probabilities for pronouns with occupations.

[ArjunSubramonian]

I think @epwalsh or @dirkgr may be more equipped to answer this question.

[epwalsh]

Hi @alicesaunders, one thing you could do is explicitly set tell the batch sampler which sorting keys to use. For example, in your config:

"data_loader": {
  "batch_sampler": {
    "type": "bucket",
    "sorting_keys": ["tokens"],
  }
}

The strings in the sorting_keys array should correspond to the names of fields in your instances.

[epwalsh]

But it's still odd that you got that error. Which dataset reader are you using? Or, if you're using a custom dataset reader, what types of fields does your dataset reader put into Instance objects?

[alicesaunders]

It ended up being quite an obvious error in the config file - I had only partially filled out the information needed for the coref dataset reader needed in the config file.

I've run into another issue down the line though. I'm using Colab Pro (which should be 25MB of RAM), however, when training the spanBERT-large model for coreference resolution using a customised coref predictor with an additional line that would return coreference probabilities (outlined above: coreference_probs = util.masked_softmax(coreference_scores, top_span_mask.unsqueeze(-1)) output_dict["probs"] = coreference_probs) I kept running out of memory. I instead used spanBERT-base (also decreasing max_sentences, max_antecedents, and spans_per_word to reduce memory).

I'm finding now when trying to use the AdversarialBiasMitigator on my trained model, the following error arises:

2021-09-11 15:52:50,762 - CRITICAL - root - Uncaught exception
Traceback (most recent call last):
File "/usr/local/bin/allennlp", line 8, in
sys.exit(run())
File "/usr/local/lib/python3.7/dist-packages/allennlp/main.py", line 46, in run
main(prog="allennlp")
File "/usr/local/lib/python3.7/dist-packages/allennlp/commands/init.py", line 122, in main
args.func(args)
File "/usr/local/lib/python3.7/dist-packages/allennlp/commands/train.py", line 121, in train_model_from_args
file_friendly_logging=args.file_friendly_logging,
File "/usr/local/lib/python3.7/dist-packages/allennlp/commands/train.py", line 187, in train_model_from_file
return_model=return_model,
File "/usr/local/lib/python3.7/dist-packages/allennlp/commands/train.py", line 260, in train_model
file_friendly_logging=file_friendly_logging,
File "/usr/local/lib/python3.7/dist-packages/allennlp/commands/train.py", line 504, in _train_worker
metrics = train_loop.run()
File "/usr/local/lib/python3.7/dist-packages/allennlp/commands/train.py", line 577, in run
return self.trainer.train()
File "/usr/local/lib/python3.7/dist-packages/allennlp/training/gradient_descent_trainer.py", line 750, in train
metrics, epoch = self._try_train()
File "/usr/local/lib/python3.7/dist-packages/allennlp/training/gradient_descent_trainer.py", line 773, in _try_train
train_metrics = self._train_epoch(epoch)
File "/usr/local/lib/python3.7/dist-packages/allennlp/training/gradient_descent_trainer.py", line 490, in _train_epoch
batch_outputs = self.batch_outputs(batch, for_training=True)
File "/usr/local/lib/python3.7/dist-packages/allennlp/training/gradient_descent_trainer.py", line 383, in batch_outputs
output_dict = self._pytorch_model(**batch)
File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1071, in _call_impl
result = forward_call(*input, **kwargs)
File "/usr/local/lib/python3.7/dist-packages/allennlp/fairness/adversarial_bias_mitigator.py", line 121, in forward
predictor_output_dict = self.predictor.forward(*args, **kwargs)
File "/content/coref_adv.py", line 231, in forward
span_mention_scores, span_mask, num_spans_to_keep
File "/usr/local/lib/python3.7/dist-packages/allennlp/nn/util.py", line 2043, in masked_topk
fill_value, _ = top_indices.max(dim=1, keepdim=True)
IndexError: max(): Expected reduction dim 1 to have non-zero size.

I'm not exactly sure what the error means, but it sounds as though the coref-spanbert-base model I have trained didn't train properly (at least that's my interpretation)? Do you know why this might be?

[alicesaunders]

I think the error is due to my output_dict not being converted correctly into a tensor or selecting the important words (which was outlined below).

coreference_probs is what will get passed into your adversary. Thus, you need to prepend a step to feedforward_regression_adversary (implemented here that selects the important coreference probabilities (the ones related to pronouns coreferring with occupations) and form them into a tensor. (In the guide example, this was not necessary because the predictor, i.e. RoBERTA for textual entailment, by default output a tensor of 3 probabilities (for entails, contradicts, and neutral) in its output_dict.)

@ArjunSubramonian how would you implement this step? I can't work out how to select the important coreference probabilities. so far I have added in a step here that looks like: self.predictor_output_key = torch.Tensor(predictor_output_dict[self.predictor_output_key]) but i think this isn't quite right...
thank you for your help!