Key and Query are multiplied in an opposite direction

[baehyunsol]

Hello, I love your implementation, thanks so much!! I've been trying to implement KV cache in femtoGPT, and I found something weird. I'd like to share it with you and get some information.

In the implementation, it transposes q matrix and multiplies k matrix with q^t (k * q^t). Below is the code

femtoGPT/src/gpt.rs

Lines 206 to 207 in f0afe9e

	let q_t = g.call(Transpose::new(), &[q])?;
	let kq = g.call(MatMul::new(), &[k, q_t])?;

But in the paper, it is q * k^t.

Also, in nanoGPT, it's q * k^t. code

---

As far as I know, using k * q^t instead of q * k^t doesn't degrade the performance of the model. K in femtoGPT does what Q in the other GPT does, and Q in femtoGPT does what K in other GPT does. It's like using the name oppositely.

There's a dilemma:

1. If we don't fix it, femtoGPT and the other GPTs use names oppositely. It might confuse the new comers.
2. If we do fix it, all the previously trained models are not compatible with the new version of femtoGPT. We can manually swap head_{}_{}_q tensors and head_{}_{}_k tensors, but that'd be a lot of work.

[baehyunsol]

I did a very simple experiment. I trained 2 models, one with k * q^t (original femtoGPT), and another with q * k^t. I used this code to generate a synthetic data. Both models have embedding_degree 96, num_layers 6, num_heads 4 and num_tokens 64.

• test 1: k * q^t (original femtoGPT)
  • initial loss 3.784
  • loss at step 50: 3.716
  • loss at step 100: 3.332
  • loss at step 150: 3.227
  • loss at step 200: 3.057
  • loss at step 250: 2.989
• test 2: q * k^t
  • initial loss 3.784
  • loss at step 50: 3.715
  • loss at step 100: 3.329
  • loss at step 150: 3.170
  • loss at step 200: 2.741
  • loss at step 250: 2.671
  • loss at step 300: 2.554
  • loss at step 350: 2.291

There's supposed to be no difference in the model performance. There is a difference but I guess that's a coincidence. I manually evaluated the outputs of both models, and they had a similar quality.

[baehyunsol]

I did another experiment.

1. I fixed k * q^t to q * k^t. I'll call it new version of femtoGPT.
2. I have many checkpoints that are trained with the original version of femtoGPT (k * q^t). When I run them with the new version of femtoGPT, it works but the models generate nonsense output.
3. If I swap head_{}_{}_k with head_{}_{}_q of the models, they generate good output. I also checked their losses, and the losses are identical to the old losses.

So it confirms that this issue has nothing to do with the model's performance. It only has to do with swapping head_{}_{}_k with head_{}_{}_q and breaking backward compatibilities.

[keyvank]

Thanks for sharing! It's ok i think let's switch them