Clarify shape descriptions inside forward method #9
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
I think that shape descriptions in forward method are somewhat misleading.
The input has size (B, T, C) and then this
C
is used to describe shape on every step during forward pass, where in fact the last dimension is actuallyhead_size
.C
can be equal tohead_size
(if for example the number of heads is 1), but it's not guaranteed.In addition the shape of projection layer should have
(head_size * num_heads, n_embd)
as input and output dimensions respectively. This will make it more robust. In this code the size of head is equal ton_embd // num_heads
, but it's not a strict rule and we might want to not reduce the size of a single head (by a factor of num_heads) and as a result the last dimension of concatenated heads will be larger than embedding size. Like it's shown in this article.