Exercise on Transformers --- Answer to problem 1