Make MultiHeadAttention use masks from query and value tensors #7951
Conversation
tfjs-layers/src/engine/topology.ts
Outdated
| for (let i = 0; i < output.length; i++) { | ||
| output[i].kerasMask = outputMask[i]; | ||
| } | ||
| } else if (outputMask instanceof Array) { |
There was a problem hiding this comment.
what if the array only contains one mask?
There was a problem hiding this comment.
I think that should be an error. If there's only one mask for all the tensors, it should be returned as a tensor instead of a [tensor].
There was a problem hiding this comment.
Actually, keras does not seem to broadcast masks at all. Each output tensor needs its own mask:
https://github.com/keras-team/keras/blob/master/keras/engine/base_layer.py#L2893-L2898
| if (useCausalMask) { | ||
| // the shape of the causal mask is [1, T, S] | ||
| const mask = this.computeCausalMask(query, value); | ||
| autoMask = mask; | ||
| autoMask = autoMask ? logicalAnd(autoMask, mask) : mask; |
There was a problem hiding this comment.
how is this associated with the Topology computeMask logic?
There was a problem hiding this comment.
Earlier layers (Embedding) have computeMask called to compute the mask for their output tensors. This layer uses those masks.
mattsoulanille
force-pushed
the
fix_mha
branch
from
111048c to
6822278
Compare
| type MaybeSymbolic = SymbolicTensor | Tensor; | ||
|
|
||
| function checkAllSymbolic(tensors: MaybeSymbolic | MaybeSymbolic[] | ||
| ): tensors is SymbolicTensor | SymbolicTensor[] { | ||
| let allAreSymbolic = true; | ||
| for (const tensor of generic_utils.toList(tensors)) { | ||
| if (!(tensor instanceof SymbolicTensor)) { | ||
| allAreSymbolic = false; | ||
| break; | ||
| } | ||
| } | ||
| return allAreSymbolic; | ||
| } | ||
|
|
||
| function checkNoneSymbolic(tensors: MaybeSymbolic | MaybeSymbolic[] | ||
| ): tensors is Tensor | Tensor[] { | ||
| let noneAreSymbolic = true; | ||
| for (const tensor of generic_utils.toList(tensors)) { | ||
| if (tensor instanceof SymbolicTensor) { | ||
| noneAreSymbolic = false; | ||
| break; | ||
| } | ||
| } | ||
| return noneAreSymbolic; | ||
| } |
There was a problem hiding this comment.
I moved this here so I could write them as type guards (tensors is Tensor...)
| @@ -188,101 +188,6 @@ describe('MultiHeadAttention', () => { | |||
| expectTensorsNotClose(queryKernel, outputKernel, 1e-6); | |||
| }); | |||
|
|
|||
| describeMathCPU('High Dimensional Attention', () => { | |||
| /** | ||
| * Test that the value and causal masks are taken into account. | ||
| */ | ||
| function testValueMask(testcaseName: string, useCausalMask: boolean) { |
There was a problem hiding this comment.
Refactored to explicitly declare both tests instead of using a loop to declare them.
| @@ -482,6 +395,101 @@ describe('MultiHeadAttention', () => { | |||
| // TODO(pforderique): Test serialization. | |||
| }); | |||
|
|
|||
| describeMathCPU('High Dimensional Attention', () => { | |||
There was a problem hiding this comment.
Moved some tests from above to here.
mattsoulanille
requested review from
fengwuyao
and removed request for
Linchenn
Add optional Keras masks to tfjs tensors. Enable them for tfjs-layers on layers that emit them. Use the masks of query and value input tensors in MultiHeadAttention to compute the correct mask automatically.
To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.