Transformer architecture not working on toy problem My transformer is not working on a toy problem.

Toy problem
Input : Sequence of random integer, one-hot-encoded. Example :
[[0, 0, 0, 1, 0, 0, 0, 0, 0, 0],
 [0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
 [0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
 [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
 [0, 0, 0, 1, 0, 0, 0, 0, 0, 0]
]

Output : 1 if the first random integer is &lt; 5, 0 otherwise. For the given example, the output would be 1.

Problem
Running my Transformer on this toy problem works perfectly. However, running with a bigger size (one-hot encoding on 800 dimension) does not work anymore : the network always output similar results, no matter the input.
Note : I did change the condition of the output to be x &lt; 800, in order to have class balance.
Why my transformer architecture is not working with this variation of the toy problem ?

Another problem
I also tried to use 2 inputs (similarly generated, with size = 10) and generate the output based on these 2 inputs (1 if the first random integer of both input is &lt; 5 or both &gt; 5, 0 otherwise)
I am not sure but maybe the excess zeroes are the culprit. I guess the results you are getting are just the accumulation of the biases. I would suggest converting your one-hot vectors to dense vectors using an Embedding Layer.