Vision cross attention mask transform

[RdoubleA]

Context

Multimodal models that use cross-attention layers in the text LLM backbone to attend to outputs of a vision encoder on the images require a cross-attention mask to ensure the correct text tokens attend to the right images. For multiple, interleaved images in a single sample, we follow the approach in Flamingo (Fig. 7 of https://arxiv.org/pdf/2204.14198) where text tokens after an image attend to that particular image (more details in docstrings). To create this mask, this PR adds a transform class that takes a sample dictionary (i.e., during data preprocessing in one of torchtune's dataset classes), computes this cross-attention mask and adds it to the dictionary.

Handling variable number of tiles and number of images across samples will be done in the batch collator in a follow-up PR. For now, we return the masks as a list of tensors, one mask per image, which can have varied n_tiles.

Changelog

• Add base Transform interface
• Add VisionCrossAttentionMask
• Add unit tests for VisionCrossAttentionMask

Test plan

pytest tests/torchtune/modules/transforms/test_transforms.py

Docs

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1141

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit bead59a with merge base 37636a8 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[felipemello1]

it looks great! I left a few comments more around naming, since tiles/patches/image is so confusing in general. I think that this class would benefit from a visual small example too.

regarding the interval parts, IMO we should try to make the operations as explicit as possible. Instead of tok1, tok2; i; vision_mask[0], we should try to give them names, like idx_start, idx_end, etc.

[felipemello1]

can we make it more visual? Something like this: https://github.com/huggingface/transformers/blob/60bb571e993b7d73257fb64044726b569fef9403/src/transformers/models/llava_next/modeling_llava_next.py#L446

Or a link to the paper + page where they have an image for it

[felipemello1]

if this CrossAttention is specific for text + images, should we indicate it in the name? Something like MultimodalCrossAttentionMask or VisionTextCrossAttentionMask?

[RdoubleA]

Yeah that's a good point. Will rename

[felipemello1]

I wonder if we should add a link to modules/VisionTransformer for a in-depth explanation of what num_patches mean. For better clarity, would it make sense to rename it num_patches_per_tile, since later we multiply it by n_tiles?

If we say "number of patches per image", it may be confusing, because an image can have a variable number of patches.

later on you say:

single_image_seq_len = n_tiles * self.num_patches + 1
image_seq_len = single_image_seq_len * n_img

So tile != image. Image is a set of tiles.

[RdoubleA]

ah sorry, probably my shallow understanding of patches/images/tiles. What I intended was num_patches per tile. If it makes sense I'd like to keep the name consistent with your vision transformer (either patch_grid_size or patch_size maybe?), whichever parameter you use to compute num patches.

I also assumed at this point tiles is padded to the max in all the images. Is this incorrect? Where does the padding happen?

[felipemello1]

patches_per_tile is a fixed size, and its calculated as (tile_size // patch_size)**2.

What I did in VisionTransformer was to ask the user to pass tile_size and patch_size, and I calculated it for them. The VisionTransformer has a helper function that saves this value: https://github.com/felipemello1/torchtune/blob/f683812626ad4559464840112ddce516487bea5c/torchtune/modules/vision_transformer.py#L249

Maybe get it from the model, or ask for tile_size and patch_size, to avoid user confusion?

[felipemello1]

Since the torchvision compose has a different behavior, I wonder if it makes sense to change Compose to something else, so users dont get confused with tv.Compose. Maybe "ComposeTransforms"?

[RdoubleA]

how about Pipeline?

[ebsmothers]

Just for my own understanding, the main difference with torchvision compose is that we support multiple inputs and multiple outputs here? Can we not just use torchvision compose with a single dict?

[felipemello1]

I tried naming something as Pipeline, as Kartikay said it would confuse people, because it is also used by other libraries :P. I guess sklearn?

[RdoubleA]

@ebsmothers our Compose needs to have a slightly different forward signature to unfold dictionary inputs. From torchvision:

def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img

but to avoid confusion, I agree should name it something else. Just haven't figured out what yet

[felipemello1]

nit: should we add Args:,Returns, Examples?

[felipemello1]

nit: Maybe add a comment with the other dimensions, so we know what they are, but keep the "_", so we know they are not used?

You said "# We are still at sample level pre-collating"
So is n_img == bsz? If so, for consistency with VisionTransformer, should we rename it?

[ebsmothers]

yeah please add type and shape info for arguments to __call__

[felipemello1]

nit: maybe add comment explaining that +1 is for CLS, if thats the case

[ebsmothers]

nit: split line differently if linter allows it, as written this is confusing

[ebsmothers]

Why do we need to also update with tokens and images? Isn't this a no-op for those args?

[RdoubleA]

Since they are explicit keyword args, they get unfolded from kwargs and you have to add them back in

[ebsmothers]

Maybe it's just me but I find this whole kwargs pattern kinda unintuitive. Like really we are covertly using this transform to also add its inputs into the dictionary (does that mean that it doesn't already contain them?). Maybe I need to see the callsite, but I don't understand why we can't just have e.g.

sample: Dict[str, Any]
mask_transform = VisionCrossAttentionMask(...)
sample["encoder_mask"] = mask_transform(sample["tokens"], sample["images"])

Alternatively if you want to go more of the compose route:

class VisionCrossAttentionMask:
	def __call__(sample: Dict[str, Any]) -> Dict[str, Any]:
		# just key into sample directly in the transform
		sample["encoder_mask"] = construct_mask(sample["tokens"], sample["images"])
		return sample

mask_transform = VisionCrossAttentionMask(...)
sample = mask_transform(sample)

But rn we are kinda halfway in between the two and it feels clunky to me.

[RdoubleA]

cc @pbontrager who'sthought about this a lot more

We were between two approaches, the kwargs approach here and the sample approach you described where you key into the fields you need for the transform. One issue with that was this implicit assumption that token ids are always under "tokens", image tensors are always under "images", etc., so we thought we may have to provide these keys as attributes to the transform class. Then the key management gets a bit messy to handle and error prone.

Although, thinking about this more, the collator will also implicitly assume the keys "tokens", "labels", "images" etc will be present. If we're consistent with these across the library, and the fact that the number of these keys in the sample dict will remain quite small ~5-10 max, it might be fine to keep these assumptions, or do something similar to checkpointing with model key, optim key, etc as constants.

I don't have a strong opinion either way, although it would be nice to keep a forward signature of just sample.

[ebsmothers]

Yeah imo you're gonna need the keys somewhere anyways, better to just be explicit about it than do this weird back and forth between dict values and standalone arguments.

[felipemello1]

mostly nits about naming/docs. Feel free to ignore.

[felipemello1]

.