Extension of WKV from 1D to 2D.

[Yaziwel]

It is an impressive extension of WKV from unidirectional to bidirectional. However, the current WKV is still limited to 1D sequences. I believe the relative bias |t-i|/T can be extended to (|x-i|/H + |y-j|/W), where x and y represent two directions in an image. It appears that this approach can also be transformed into a recurrent form with linear complexity.

[duanduanduanyuchen]

It is an impressive extension of WKV from unidirectional to bidirectional. However, the current WKV is still limited to 1D sequences. I believe the relative bias |t-i|/T can be extended to (|x-i|/H + |y-j|/W), where x and y represent two directions in an image. It appears that this approach can also be transformed into a recurrent form with linear complexity.

You're right, Bi-WKV in 2D can still maintain linear complexity. We considered this implementation when we designed vrwkv, but realized that 1D was good enough and did not modify it to a 2D version to be more closely related to the original rwkv (make rwkv applicable to visual tasks with only fewer modifications). 2D WKV may be a better structure in visual tasks.

[Yaziwel]

It is an impressive extension of WKV from unidirectional to bidirectional. However, the current WKV is still limited to 1D sequences. I believe the relative bias |t-i|/T can be extended to (|x-i|/H + |y-j|/W), where x and y represent two directions in an image. It appears that this approach can also be transformed into a recurrent form with linear complexity.

You're right, Bi-WKV in 2D can still maintain linear complexity. We considered this implementation when we designed vrwkv, but realized that 1D was good enough and did not modify it to a 2D version to be more closely related to the original rwkv (make rwkv applicable to visual tasks with only fewer modifications). 2D WKV may be a better structure in visual tasks.

Do you have plans to rewrite this part of the CUDA kernel to extend it to 2D? I believe many people, including myself, need this but don't know how to write CUDA.

[duanduanduanyuchen]

It is an impressive extension of WKV from unidirectional to bidirectional. However, the current WKV is still limited to 1D sequences. I believe the relative bias |t-i|/T can be extended to (|x-i|/H + |y-j|/W), where x and y represent two directions in an image. It appears that this approach can also be transformed into a recurrent form with linear complexity.

You're right, Bi-WKV in 2D can still maintain linear complexity. We considered this implementation when we designed vrwkv, but realized that 1D was good enough and did not modify it to a 2D version to be more closely related to the original rwkv (make rwkv applicable to visual tasks with only fewer modifications). 2D WKV may be a better structure in visual tasks.

Do you have plans to rewrite this part of the CUDA kernel to extend it to 2D? I believe many people, including myself, need this but don't know how to write CUDA.

Thanks for your advice! I will add this idea to my schedule. Since I'm still busy with other work these days, I will reply once I finish.

[duanduanduanyuchen]

It is an impressive extension of WKV from unidirectional to bidirectional. However, the current WKV is still limited to 1D sequences. I believe the relative bias |t-i|/T can be extended to (|x-i|/H + |y-j|/W), where x and y represent two directions in an image. It appears that this approach can also be transformed into a recurrent form with linear complexity.

You're right, Bi-WKV in 2D can still maintain linear complexity. We considered this implementation when we designed vrwkv, but realized that 1D was good enough and did not modify it to a 2D version to be more closely related to the original rwkv (make rwkv applicable to visual tasks with only fewer modifications). 2D WKV may be a better structure in visual tasks.

Do you have plans to rewrite this part of the CUDA kernel to extend it to 2D? I believe many people, including myself, need this but don't know how to write CUDA.

Thanks for your advice! I will add this idea to my schedule. Since I'm still busy with other work these days, I will reply once I finish.

[chenzean]

yes,超级期待作者可以将其扩展到2D！！！！

[Yaziwel]

It is an impressive extension of WKV from unidirectional to bidirectional. However, the current WKV is still limited to 1D sequences. I believe the relative bias |t-i|/T can be extended to (|x-i|/H + |y-j|/W), where x and y represent two directions in an image. It appears that this approach can also be transformed into a recurrent form with linear complexity.

You're right, Bi-WKV in 2D can still maintain linear complexity. We considered this implementation when we designed vrwkv, but realized that 1D was good enough and did not modify it to a 2D version to be more closely related to the original rwkv (make rwkv applicable to visual tasks with only fewer modifications). 2D WKV may be a better structure in visual tasks.

Do you have plans to rewrite this part of the CUDA kernel to extend it to 2D? I believe many people, including myself, need this but don't know how to write CUDA.

Thanks for your advice! I will add this idea to my schedule. Since I'm still busy with other work these days, I will reply once I finish.

I hope this extension can be based on RWKV4 because I found that the training speed of RWKV6 is too slow for my task. Considering this, RWKV-4 might be more appealing.