Fixed Masking in HookedTransformer.generate

[tuomaso]

Description

Currently calling HookedTransformer.generate() with a batch of inputs, i.e. a list of two prompts does not correctly mask the padding tokens, causing the batched outputs to be different from the outputs of generating 1-by-1 for the shorter sequences. Current code does not mask attention at all, so I implemented a ~5 line change to mask attention correctly using existing functions. This fix only works when padding_side="left" so I changed the default value for generate to be "left", not sure how to fix this for right padding. After the fix .generate() with a batch gives the same output as generating one by one (with do_sample = False).

Fixes # (issue)

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Changing the default value/fixing behavior will change some outputs but only ones that were broken already.

Screenshots

Before

Please attach before and after screenshots of the change if applicable.

After

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[danwilhelm]

I've run into this problem as well. Due to no attention masking, the TransformerLens generate returns different results than Transformers generate when passed >1 prompts. The provided changes here remedy this issue in my case.

[danwilhelm]

It may be best to update run_with_hooks and run_with_cache as well, since their >1 prompt output is now inconsistent with generate.

However, I suspect this proposed solution should not be made default. Existing users likely expect the existing behavior, even though it is inconsistent with Transformers. Perhaps it would be best to enable the attention masking behavior via a flag instead?