Transformer architecture not working on toy problem

Question

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 < 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 < 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 < 5 or both > 5, 0 otherwise)

Answer 1

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.