(Prototype) MaskedTensor Overview¶
This tutorial is designed to serve as a starting point for using MaskedTensors and discuss its masking semantics.
MaskedTensor serves as an extension to
torch.Tensor that provides the user with the ability to:
use any masked semantics (for example, variable length tensors, nan* operators, etc.)
differentiation between 0 and NaN gradients
various sparse applications (see tutorial below)
For a more detailed introduction on what MaskedTensors are, please find the torch.masked documentation.
In this section we discuss how to use MaskedTensor including how to construct, access, the data and mask, as well as indexing and slicing.
We’ll begin by doing the necessary setup for the tutorial:
import torch from torch.masked import masked_tensor, as_masked_tensor import warnings # Disable prototype warnings and such warnings.filterwarnings(action='ignore', category=UserWarning)
There are a few different ways to construct a MaskedTensor:
The first way is to directly invoke the MaskedTensor class
Throughout this tutorial, we will be assuming the import line: from torch.masked import masked_tensor.
Accessing the data and mask¶
The underlying fields in a MaskedTensor can be accessed through:
MaskedTensor.get_mask()function. Recall that
Trueindicates “specified” or “valid” while
Falseindicates “unspecified” or “invalid”.
In general, the underlying data that is returned may not be valid in the unspecified entries, so we recommend that
when users require a Tensor without any masked entries, that they use
MaskedTensor.to_tensor() (as shown above) to
return a Tensor with filled values.
Indexing and slicing¶
MaskedTensor is a Tensor subclass, which means that it inherits the same semantics for indexing and slicing
torch.Tensor. Below are some examples of common indexing and slicing patterns:
data: tensor([[[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]]) mask: tensor([[[ True, False, True, False], [ True, False, True, False], [ True, False, True, False]], [[ True, False, True, False], [ True, False, True, False], [ True, False, True, False]]])
mt: MaskedTensor( [ [ 0.0000, --, 2.0000, --], [ 4.0000, --, 6.0000, --], [ 8.0000, --, 10.0000, --] ] ) mt[:, :, 2:4]: MaskedTensor( [ [ [ 2.0000, --], [ 6.0000, --], [ 10.0000, --] ], [ [ 14.0000, --], [ 18.0000, --], [ 22.0000, --] ] ] )
Why is MaskedTensor useful?¶
MaskedTensor’s treatment of specified and unspecified values as a first-class citizen
instead of an afterthought (with filled values, nans, etc.), it is able to solve for several of the shortcomings
that regular Tensors are unable to; indeed,
MaskedTensor was born in a large part due to these recurring issues.
Below, we will discuss some of the most common issues that are still unresolved in PyTorch today
and illustrate how
MaskedTensor can solve these problems.
Distinguishing between 0 and NaN gradient¶
One issue that
torch.Tensor runs into is the inability to distinguish between gradients that are
undefined (NaN) vs. gradients that are actually 0. Because PyTorch does not have a way of marking a value
as specified/valid vs. unspecified/invalid, it is forced to rely on NaN or 0 (depending on the use case), leading
to unreliable semantics since many operations aren’t meant to handle NaN values properly. What is even more confusing
is that sometimes depending on the order of operations, the gradient could vary (for example, depending on how early
in the chain of operations a NaN value manifests).
MaskedTensor is the perfect solution for this!
In Issue 10729, we notice a case where the order of operations
can matter when using
torch.where() because we have trouble differentiating between if the 0 is a real 0
or one from undefined gradients. Therefore, we remain consistent and mask out the results:
tensor([4.5400e-05, 6.7379e-03, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, nan, nan])
MaskedTensor( [ 0.0000, 0.0067, --, --, --, --, --, --, --, --, --] )
The gradient here is only provided to the selected subset. Effectively, this changes the gradient of where to mask out elements instead of setting them to zero.
Issue 52248 is another example.
torch.where(b, a/0, c): tensor(1., grad_fn=<WhereBackward0>) torch.autograd.grad(torch.where(b, a/0, c), a): (tensor(nan),)
torch.where(b, a/0, c): MaskedTensor( 1.0000, True) torch.autograd.grad(torch.where(b, a/0, c), a): (MaskedTensor(--, False),)
This issue is similar (and even links to the next issue below) in that it expresses frustration with unexpected behavior because of the inability to differentiate “no gradient” vs “zero gradient”, which in turn makes working with other ops difficult to reason about.
When using mask, x/0 yields NaN grad¶
In Issue 4132, the user proposes that
x.grad should be [0, 1] instead of the [nan, 1],
MaskedTensor makes this very clear by masking out the gradient altogether.
MaskedTensor( [ --, 1.0000] )
In Issue 67180,
the gradient isn’t calculate properly (a longstanding issue), whereas
MaskedTensor handles it correctly.
MaskedTensor( 3.0000, True)
Safe softmax is another great example of an issue that arises frequently. In a nutshell, if there is an entire batch that is “masked out” or consists entirely of padding (which, in the softmax case, translates to being set -inf), then this will result in NaNs, which can lead to training divergence.
MaskedTensor has solved this issue. Consider this setup:
x: tensor([[ 0.2345, -inf, -inf], [-0.1863, -inf, -0.6380], [ -inf, -inf, -inf]]) mt: MaskedTensor( [ [ 0.2345, --, --], [ -0.1863, --, -0.6380], [ --, --, --] ] )
For example, we want to calculate the softmax along dim=0. Note that the second column is “unsafe” (i.e. entirely masked out), so when the softmax is calculated, the result will yield 0/0 = nan since exp(-inf) = 0. However, what we would really like is for the gradients to be masked out since they are unspecified and would be invalid for training.
tensor([[0.6037, nan, 0.0000], [0.3963, nan, 1.0000], [0.0000, nan, 0.0000]])
MaskedTensor( [ [ 0.6037, --, --], [ 0.3963, --, 1.0000], [ --, --, --] ] )
Implementing missing torch.nan* operators¶
In Issue 61474,
there is a request to add additional operators to cover the various torch.nan* applications,
In general, these problems lend themselves more naturally to masked semantics, so instead of introducing additional
operators, we propose using
Since nanmean has already landed,
we can use it as a comparison point:
y: tensor([ 0., 1., 4., 9., 0., 5., 12., 21., 0., 9., 20., 33., 0., 13., 28., 45.]) z: tensor([nan, 1., 4., 9., nan, 5., 12., 21., nan, 9., 20., 33., nan, 13., 28., 45.])
y.mean(): tensor(12.5000) z.nanmean(): tensor(16.6667) torch.mean(masked_tensor(y, y != 0)): MaskedTensor( 16.6667, True)
In the above example, we’ve constructed a y and would like to calculate the mean of the series while ignoring
the zeros. torch.nanmean can be used to do this, but we don’t have implementations for the rest of the
MaskedTensor solves this issue by being able to use the base operation,
and we already have support for the other operations listed in the issue. For example:
MaskedTensor( 1.0000, True)
Indeed, the index of the minimum argument when ignoring the 0’s is the 1 in index 1.
MaskedTensor can also support reductions when the data is fully masked out, which is equivalent
to the case above when the data Tensor is completely
nanmean would return
(an ambiguous return value), while MaskedTensor would more accurately indicate a masked out result.
x: tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan]) torch.nanmean(x): tensor(nan) torch.nanmean via maskedtensor: MaskedTensor(--, False)
This is a similar problem to safe softmax where 0/0 = nan when what we really want is an undefined value.
In this tutorial, we’ve introduced what MaskedTensors are, demonstrated how to use them, and motivated their value through a series of examples and issues that they’ve helped resolve.