[Python] Difference between tokens generated on a configuration in two different contexts

I have a model that given a configuration, or state (of a Rubik's cube, but whatever, it is a sequence of integers) generates a movement (from 0 to 5). This movement can be used to bring the configuration to another state and the model has been initially trained to bring the configuration to a solution, an identical state, by generating in a cycle every token, either a movement or a number in the configuration, or a number encoding a space. On the other hand, we also can rearrange a state in a deterministic way by means of a function called reindex, below. Now, I can generate plenty of configurations with a function that is called bould_state and run the following:

def reindex(state,moves):
#args are tensors
for m in dict_generators[moves]:
state = state[m]
return(state)

def build_state(n):
seq_moves = torch.randint(0, 5, (n,), device=device)
state = reindex(identical_state,seq_moves)
return state

for i in range(400):
#encoding
x = build_state(int(i / 20)+1)
x = x + 8
# add "next line" tokens, generate encoded torch.Size([26])
x = torch.cat((x, torch.tensor([7], device=device)), dim=0)
x = torch.cat((torch.tensor([7], device=device), x), dim=0)
print(model.generate(x[None],temperature=temperature,top_k=top_k))


As you can see, the model does not properly eat a configuration but an encoding, which is done in order to have all of the number in the configuration shifted by 8, and using the first eight tokens to either give a movement (from 0 to 5) or saying that the configuration reached the desired resolved state (6), or again to separate two configuration (7). The code above works pretty well, and it actually generates only movements token except in extremely few cases, when, based on the training it sometimes return a 8 token, which is the encoding of 0 and the beginnning of evey configuration.

Now, something I cannot understand happens. If I run a code that wants to generate not only a movement but the whole sequence for a solution, by a cycle in which I generate a movement and I reindex the same one (with respective deconding-encodig), the number of "wrong" tokens 8 becomes insanely much higher, even though what I pass to the model is just a configuration as in the previous case, only not a new one but a reindex of the previous one. For example, the code

def predict_solution(x):
#arg tensor torch.Size([24])
sol = []
list_probs = []
# encode:
x = x + 8
# add next line tokens, generate encoded torch.Size([26])
x = torch.cat((x, torch.tensor([7], device=device)), dim=0)
x = torch.cat((torch.tensor([7], device=device), x), dim=0)
#make it torch.Size([1,26]) to pass to the model
x = x[None]
for it in range(max_sol_length):
y,probs = model.generate(x, temperature=temperature, top_k=top_k)
#y, probs are two gradients os Size [1,1], probs keeps gradient

if torch.equal(y torch.tensor([[8]]))
print("wrong token")
#force it someway to be a movement in order to apply reindex below

list_probs.append(probs)
sol.append(y)

#decode x into a tensor of torch.Size([24]) to pass to reindex
x = x.squeeze(0)
x = x-8
x = x[1:-1]
x = reindex(x,y)
x = x[None] #unsqueeze(0)

if torch.equal(x, identical_state[None]):
return sol, list_probs

#re-encode to pass to the model again
x = x + 8
x = torch.cat((x, torch.tensor([[7]], device=device)), dim=1)
x = torch.cat((torch.tensor([[7]], device=device), x), dim=1)

if it == max_sol_length - 1:
print("solution not found")
return sol, list_probs


for i in range(400):
x = build_state(int(i / 20)+1)
predict_solution(x)


It seems like ar each cycle the argument of generate remembers that we are not providing it with a new freshly generated tensor, but with a reindexed. The thing to know is also that I would like to keep the gradient for the probabilities sampled by the model together with the movement in order to keep doing some reinforcement learning on that model. Why does that happen and how to fix it? Is there a way to do that without going deep in re-training the original model?

I tried to also use a detach() to x before passing it to the model, at the end of the cycle (even though this clearly would break the computational graph, so it would not be an admissible solution to my problem), but the problem persists!

Continue reading...