Self documenting Payload instead of Tuples as output of Transformer
Author: bhoov
Pull request: no
Labels: name: Core: Modeling

I propose a replacing the default Tuple outputs with payloads where every field in the tuple is accessible by a name. This has the following benefits:

1. Being more accessible to newcomers of the library
2. Eliminating suspicious comments in the source code describing the output (see below)
3. Make it easier to extract particular values from different models for inspection -- adding a named field is easier than mangling the order of an existing Tuple
4. Like 3, where new models can output their own metadata without reordering expected structure
5. A consistent output structure makes it easier to compare different models in applications that involve multiple architectures.


## Motivation
This will be my first formal feature request, as I am getting weary of forgetting what the Tuple output in the forward pass of a Transformer means. I find myself constantly returning to the source code for each model and playing with the shapes and values of the different fields to check whether something is what I expect it to be.

### The Problem

Currently, the output of a Transformer (I use Bert as an example here, potentially not the most recent version) is structured and documented in the source code as follows:

``` python
...
    for i, layer_module in enumerate(self.layer):
        if self.output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        layer_outputs = layer_module(
            hidden_states, attention_mask, head_mask[i], encoder_hidden_states, encoder_attention_mask
        )
        hidden_states = layer_outputs[0]

        if self.output_attentions:
            all_attentions = all_attentions + (layer_outputs[1],)

    # Add last layer
    if self.output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)

    outputs = (hidden_states,)
    if self.output_hidden_states:
        outputs = outputs + (all_hidden_states,)
    if self.output_attentions:
        outputs = outputs + (all_attentions,)
        if self.output_additional_info:
            outputs = outputs + (all_additional_info,)

    return outputs  # last-layer hidden state, (all hidden states), (all attentions)
```

Of course, that last comment is highly dependent on what you pass to the configuration. What if you desire all the attentions but not the hidden_states? Now, attention is at `outputs[1]` instead of `outputs[2]`. Utterly confusing.

And what if you have an architecture that has different outputs? Here is an example from the output of a T5 model.

``` python
...
    # Add last layer
    if self.output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)

    outputs = (hidden_states,)
    if use_cache is True:
        assert self.is_decoder, "`use_cache` can only be set to `True` if {} is used as a decoder".format(self)
        outputs = outputs + (present_key_value_states,)
    if self.output_hidden_states:
        outputs = outputs + (all_hidden_states,)
    if self.output_attentions:
        outputs = outputs + (all_attentions,)
    return outputs  # last-layer hidden state, (presents,) (all hidden states), (all attentions)
```

There's this new field `presents` that again confuses the order. It starts to get a bit confusing.
In addition, one of my projects the past many months has been to visually compare and interpret different Transformer models [exbert](http://exbert.net/). This means that I often want to edit the source code and extract, for example, the keys / values / head embeddings prior to the final projection into the embeddings that are passed to the next layer. Extracting these and passing them through the model is more complicated than it should be -- there are no hooks that I can use to catch arbitrary information within a module's forward pass (potentially a separate feature request, but I feel this would slow the prototyping speed of this library quite a bit), and I worry about messing with an expected order to the Tuple output. It is also really easy to forget which field in a list of 8 items is the one I want. Architectures are also incorporated quickly into Transformers (kudos!), and it would be great to know what inference information I have available for a model simply by looking at the object outputted by the forward pass.

### Possible Solutions

I would like the return object of every Transformer's forward pass to be a Payload where the information outputted is easily identified by the fields. E.g., for a LMHead Transformer:

``` python
{
logits: __,
past: __,
hidden_states: __,
attentions: __,
}
```
Where a non-LMHead Transformer would not include logits and offset indexing into the output. This also allows models like T5 to unambiguously add additional fields without compromising the structure of the tuple.

``` python
{
last_layer_hidden_state: __,
presents: __,
}
```

It would be trivial to add fields to this at the output: 

``` python
    output = {"last_hidden_state": hidden_states}
    if self.output_hidden_states:
        all_hidden_states = all_hidden_states + (hidden_states,)
        outputs['all_hidden_states'] = all_hidden_states
    
    if self.output_attentions:
        outputs['attentions'] = all_attentions

    return outputs  # Self documenting fields, no comment needed! :D
```

Payloads like this would also work for intermediate modules, though some naming convention to indicate the main output intended to be used by the next module/layer would be necessary.

Python's `namedtuples` could also be an option, though easily adding fields to this immutable structure is a bit more challenging for interpretability work. Additional alternatives could be [`namedlist`](https://pypi.org/project/namedlist/) or [`recordclass`](https://pypi.org/project/recordclass/).

## Contribution

I am willing to continue thinking of solutions and work towards this goal, but making a single PR would be both a sweeping change across library and the way each module is coded. However, I believe this change to be important for increasing the accessibility of the library for both newcomers and those who want to build applications around the different models.


Comments:
I agree with this change but what about backward compat? I think, as you said, `namedtuple` could be used.

I think the `namedtuple` would be best for both getting the desired features and also backward compatibility. But to make it backward compatible, wouldn't we just need any object that supports indexing (edit: and unpacking)? What other special features of a tuple are important to maintain for backwards compatibility?

Unfortunately, `torch.jit` does not support dict, namedtuples or other kinds of fancy outputs... just plain old tuples. See [this issue](https://github.com/pytorch/pytorch/issues/373440) for instance.

Dict support has been added recently, so we could consider switching to that with a breaking change once it lands in a stable release of PyTorch, but this would pin us on PyTorch 1.6.0 minimum... Not sure the benefit of this for documentation would be worth it when every output of every model is cleanly explained in its documentation.

I would argue that, despite clean documentation, there is a huge advantage to having self documenting payloads. Playing around with the outputs of the results in a jupyter notebook will give you auto completion of the fields, remove any ambiguity, and (if you are trying to develop an application that can interpret the output of as many different transformer models as possible) remove the need to look up the output format for every model and version of a model (e.g., `LMHead` or `NextSentencePrediction`).

I suppose if we wanted to avoid breaking changes with JIT, we can allow each model to have an optional parameter that enables annotation of the output or not. I believe we would find that it would ideally be enabled by default for a smoother user experience, and then any decorator that JITs a code could disable the annotation in favor of a regular tuple. But even if it is a flag that we have to manually set it would still enable better applications needing to support different models.

Thoughts then on making it an optional flag?

Yes @julien-c also gave me the idea of the flag, hadn't thought of it. You can check the PR linked above for a prototype of doing this while not breaking any backward compatibility.

Skimmed through the changes and it looks very nice! Would love to see something like this for all the models. Thanks so much for doing this!

I think this is now closed by $5438 .