Generic RotaryEmbedding Layer
#1180
Conversation
| from keras_nlp.tests.test_case import TestCase | ||
|
|
||
|
|
||
| class RotaryEmbeddingTest(TestCase): |
There was a problem hiding this comment.
We should definitely fill out the test case now that we are exposing this standalone. Maybe look of the SinePositionEncoding layer tests as a start.
There was a problem hiding this comment.
Yup, added tests from SinePositionEncoding layer, they all pass now for all backends
| matrix. It calculates the rotary encoding with a mix of sine and | ||
| cosine functions with geometrically increasing wavelengths. | ||
| Defined and formulated in [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864v4). | ||
| Takes as input the query and key tensors. The input must have shape |
There was a problem hiding this comment.
I think we should also allow this layer to take in shape [batch_size, sequence_length, feature_dim]. I'll leave some more comments below.
There was a problem hiding this comment.
This still need updating in the docstring.
| References: | ||
| - [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864v4) | ||
| """ | ||
|
|
||
| def __init__(self, max_wavelength=10000, **kwargs): |
There was a problem hiding this comment.
We should probably add two arguments here. sequence_axis=1, and feature_axis=-1, that users can set as desired.
There was a problem hiding this comment.
This would be cool. Added.
| References: | ||
| - [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864v4) | ||
| """ | ||
|
|
||
| def __init__(self, max_wavelength=10000, **kwargs): | ||
| super().__init__(**kwargs) | ||
| self.max_wavelength = max_wavelength |
There was a problem hiding this comment.
A few comments for below (github won't let me comment). Why is the following necessary?
cos_emb = cos_emb[:, : ops.shape(tensor)[1], :, :]
sin_emb = sin_emb[:, : ops.shape(tensor)[1], :, :]
The cos/sin embeddings should already have seq_len shape when you compute them.
Lastly, if you wanted to make _compute_cos_sin_embedding work with any number of dimensions, you would need to update it. Here's a draft of a change, but haven't tested this yet.
embedding = ops.concatenate((freqs, freqs), axis=-1)
for dim in range(len(x.shape)):
if axis != self.sequence_axis and axis != self.feature_axis:
embedding = ops.expand_dims(embedding, axis)
There was a problem hiding this comment.
Thanks for the pointer, got it working with some tweaks !
shivance
changed the title
RotaryEmbedding to Modeling Layers
shivance
changed the title
RotaryEmbedding to Modeling LayersRotaryEmbedding Layer
|
ValueError: Unsupported dtype19 for '{{node rotary_embedding/range}} = Range[Tidx=DT_HALF](rotary_embedding/range/start, rotary_embedding/range/Cast, rotary_embedding/range/delta)' with input shapes: [], [], [] and with computed input tensors: input[0] = <0>, input[1] = <32>, input[2] = <2>. |
| self.max_wavelength = max_wavelength | ||
| self.sequence_axis = sequence_axis | ||
| self.feature_axis = feature_axis | ||
| self.scaling_factor = scaling_factor |
There was a problem hiding this comment.
@mattdangerw LLaMa Rotary Embedding layers use a scaling factor, I've added that too!
|
/gcbrun |
shivance
changed the title
RotaryEmbedding LayerRotaryEmbedding Layer
Potentially this is just |
| matrix. It calculates the rotary encoding with a mix of sine and | ||
| cosine functions with geometrically increasing wavelengths. | ||
| Defined and formulated in [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://arxiv.org/abs/2104.09864v4). | ||
| Takes as input the query and key tensors. The input must have shape |
There was a problem hiding this comment.
This still need updating in the docstring.
This is persisting in any case. I've tried this and other typecasting. |
|
/gcbrun |
Fix dtypes with arange.
|
/gcbrun |
|
/gcbrun |
RotaryEmbedding is in the air. New SOTA models like Falcon, GPT-J, LLaMA are using it. I already developed RotaryEmbedding Layer to KerasNLP.
This is the perfect time to move it to layers and exposing it in our API should be a good addition to our API.
Closes #1087