Welcome to gateloop-transformer Discussions!

[lucidrains]

👋 Welcome!

We’re using Discussions as a place to connect with other members of our community. We hope that you:

• Ask questions you’re wondering about.
• Share ideas.
• Engage with other community members.
• Welcome others and are open-minded. Remember that this is a community we
  build together 💪.

To get started, comment below with an introduction of yourself and tell us about what you do with this community.

[lucidrains]

thoughts about paper?

[lucidrains]

i'm actually surprised max-head works as well as it does

is it necessary to have keys and values projected just to elementwise multiply them? why couldn't that just be one projection? so many unanswered questions here

[tobiaskatsch]

As an additional hint, we used group norm and a swish gate following the RetNet paper.
y = self.group_norm_function(y)
y = y * nn.swish(self.g_proj(x))
y = self.o_proj(y)
y = y.real

[lucidrains]

@tobiaskatsch

yes, this is sub-ln design, which is from the same authors. i've tried sub-ln on regular transformer but didn't see anything remarkable, but in the 'retention' setting may be different

i'll try it out tomorrow morning, thanks!

[lucidrains]

how would groupnorm make sense for gateloop if each head dimension is reduced to 1?

[jonmorton]

Isn't the swish/group norm applied to the attention output after the heads are merged? RetNet shows: GroupNorm(Concat(heads)).

Edit: I guess you'd have to pick a GN group size that's != num heads

& Perhaps the use of group norm here is really what allows you to use post-norm ..

[jonmorton]

In the updated readme you mention:

Also, gate loop's associative scan also is not able to train on even sequence lengths of 128

Is this at all because of how you had to implement it in pytorch? I.e., would you expect better results in jax with its native associate scan? The paper is also suspiciously missing any mention of training compute comparisons, so maybe it's a fundamental issue.

[lucidrains]

in jax there would be the same problem

this has been a long standing issue with all autoregressive formulations of linear attention and requires a CUDA kernel to avoid materializing the kv intermediate of dimensions (batch, seq, dim_qk, dim_v)

[tobiaskatsch]

Hello, I am the author. For the experiments, we used JAX parallel scan. In theory, the shape is (B, L, H, D, D), where H is the number of heads and D is the hidden state size. In practice, this will pose a memory bottleneck. We avoid this by choosing a max-head variant, e.g., (D=1), meaning that kv is of shape (B, L, H, 1, 1), coinciding with element-wise gating. The more general case was mainly included in the paper to highlight the connection to standard multi-head attention. This should have been explained more clearly and will be addressed in the upcoming version.

[lucidrains]

I see! Thank you for clarifying

I will rerun the experiments 🙏

[tobiaskatsch]

Also, from my experience, the quadratic O(l**2) formulation is not practical (for now) because using the cumulative product of magnitudes is numerically unstable and obvious fixes (like the clamping you try) is not sufficient in pracise. All our experience were conducted using the parallel scan variant.

[lucidrains]

Thank you

[jrb314]

Awesome work!
Any plan to replicate the main wikitext103 result in the paper (13.4 perplexity)?
Thanks!

[lucidrains]

i'm running experiments right now on enwik8, a dataset i'm familiar with. the experiments are all public (link in readme) and will grow throughout today and tomorrow.

feel free to conduct your own experiments and share them. i have everything setup in train.py

[lucidrains]

i'm very curious how the author was able to conduct wikitext103 on sufficient sequence lengths without the proper cuda kernel. the only way for this scheme to work is to fuse everything together, as we know from autoregressive linear attention line of work

[lucidrains]

but yea, let's see how it holds up to small scale experiments

[jrb314]

Paging @tobiaskatsch. If you are the paper's author, maybe you can shed some light?

[tobiaskatsch]

See #1 (comment)

[AlexDeMoor]

The official code has been released : https://github.com/tobiaskatsch/GateLoop

[lucidrains]

@AlexDeMoor nice, somebody check out the wikitext103 experiments, make sure there is no test leakage

[lucidrains]

@AlexDeMoor what happened with this? anybody ran it?