Well, yes, I deserved that reply! And yes the code is printing True. It's not that I disbelieved you... but something is wrong here. Investigation below, thanks

Well, yes, I deserved that reply! And yes the code is printing True. It's not that I disbelieved you... but something is wrong here. Investigation below, thanks to Claude.ai for walking me through it! In [10]: o1[0, :, :3] Out[10]: tensor([[ 0.0053, 0.0017, -0.0012], [ 0.0053, 0.0017, -0.0012], [ 0.0053, 0.0017, -0.0012], [ 0.0053, 0.0017, -0.0012], [ 0.0053, 0.0017, -0.0012], [ 0.0053, 0.0017, -0.0012]], grad_fn=<SliceBackward0>) Every token has the same attention values. I expect attention(cat, everything) to differ from attention(dog, everything), even without positional encoding. Further, the attention weights are uniform and identical for both sentences: In [46]: o1, aw1 = mha(W_q(e1), W_k(e1), W_v(e1)) In [47]: o2, aw2 = mha(W_q(e2), W_k(e2), W_v(e2)) In [48]: aw1.shape Out[48]: torch.Size([1, 6, 6]) In [49]: aw2.shape Out[49]: torch.Size([1, 6, 6]) In [50]: aw1 Out[50]: tensor([[[0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667]]], grad_fn=<MeanBackward1>) In [51]: aw2 Out[51]: tensor([[[0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667], [0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667]]], grad_fn=<MeanBackward1>) That is not expected. It's because the Linear layers are initialised with such small values. And the softmax causes a collapse. Trying random weights on a larger scale: In [52]: W_q.weight.data *= 100 W_k.weight.data *= 100 W_v.weight.data *= 100 In [55]: o1, aw1 = mha(W_q(e1), W_k(e1), W_v(e1)) In [56]: o2, aw2 = mha(W_q(e2), W_k(e2), W_v(e2)) In [57]: aw1 Out[57]: tensor([[[0.2049, 0.1606, 0.1256, 0.1095, 0.1723, 0.2270], [0.0883, 0.2047, 0.1544, 0.2776, 0.1405, 0.1345], [0.1196, 0.1719, 0.1831, 0.1541, 0.1374, 0.2339], [0.1413, 0.2399, 0.1617, 0.2056, 0.1634, 0.0880], [0.1455, 0.1432, 0.2432, 0.1239, 0.1494, 0.1948], [0.1897, 0.1817, 0.1920, 0.1478, 0.1618, 0.1270]]], grad_fn=<MeanBackward1>) In [58]: aw2 Out[58]: tensor([[[0.2049, 0.1606, 0.2270, 0.1095, 0.1723, 0.1256], [0.0883, 0.2047, 0.1345, 0.2776, 0.1405, 0.1544], [0.1897, 0.1817, 0.1270, 0.1478, 0.1618, 0.1920], [0.1413, 0.2399, 0.0880, 0.2056, 0.1634, 0.1617], [0.1455, 0.1432, 0.1948, 0.1239, 0.1494, 0.2432], [0.1196, 0.1719, 0.2339, 0.1541, 0.1374, 0.1831]]], grad_fn=<MeanBackward1>) In [60]: o1[:, :, :5] Out[60]: tensor([[[ 0.0145, 0.3128, -0.3659, -0.1884, 0.1724], [-0.2319, 0.1407, -0.6010, -0.4064, 0.4259], [-0.3231, 0.1622, -0.6351, -0.1711, 0.4014], [-0.0596, 0.2610, -0.7388, -0.2987, 0.3214], [-0.2750, 0.0676, -0.4140, -0.2024, 0.3383], [-0.1434, 0.0871, -0.3154, -0.0755, 0.3314]]], grad_fn=<SliceBackward0>) In [61]: o2[:, :, :5] Out[61]: tensor([[[ 0.0145, 0.3128, -0.3659, -0.1884, 0.1724], [-0.2319, 0.1407, -0.6010, -0.4064, 0.4259], [-0.1434, 0.0871, -0.3154, -0.0755, 0.3314], [-0.0596, 0.2610, -0.7388, -0.2987, 0.3214], [-0.2750, 0.0676, -0.4140, -0.2024, 0.3383], [-0.3231, 0.1622, -0.6351, -0.1711, 0.4014]]], grad_fn=<SliceBackward0>) In [62]: print("Matches: ", torch.allclose(o1, o2, atol=1e-6)) Matches: False