Backpropagation

[toth-adam]

Hi!

I have been studying your TensorFlow implementation of NEC for a while now. There is a part in your code which I do not fully understand. (Probably in the paper as well.)

The original paper says: "Backpropagation updates the the weights and biases of the convolutional embedding network and the keys and values of each action-specific memory using gradients of this loss, ..."

In your implementation, how did you solve to update the keys and values in your DND during backpropagation/backward pass? (I am curious about the DND updates especially, I think I understand how the gradients are calculated on the keys and values, I just don't see how they are get applied during backpropagation.)

Thanks in advance.

Regards,

Adam

[EndingCredits]

If I understand the paper correctly, they don't differentiate through the DND (I can't say for sure either way though). The keys in the DND are not updated through the training step (only added at the end of each episode). In the original implementation, they do update the values of saved keys if added keys are the same (they also store state hashes to check this explicitly), but I've omitted this as it was too complicated to implement, and my version (which is checking if the keys are close enough) didn't seem to really have much impact.

Updating old values is good as otherwise you're taking the mean of all samples across all time, and it takes longer and longer for updated samples to have an impact. (This is where the alpha term is used in tabular reinforcement learning). I'd be interested to hear more solutions to this problem. One idea is to keep two DNDs, one for policy, and one for learning. The agent would act according to the policy dict, and then update the learning dict, then, after a certain amount of time, the learning dict would become the new policy dict, and a new learning dict would be created.

What does happen is that every training step the networks learns to move the embedding towards the optimal position relative to stored embeddings, to align the DND predicted value and sampled value. This can then be differentiated all the way through the network. This only affects the embedding network and not the stored keys. The reason this is done is to try and obtain a good representation for the keys, but it may have the added benefit of continually pushing the embedding from older estimates (which sort of achieves the effect I describe above).

EDIT: A way to think about it is that it's basically MFEC, but the way they blend their saved values based on the sample key/embedding allows them to backprop through the network.

I hope that answers your question.

Best wishes,
Will

EDIT: Re-reading your question, "Backpropagation updates [...] the keys and values of each action-specific memory using gradients of this loss, ..." does seem to suggest they somehow update the stored keys, which I don't really see how they could do this.

I have previously spoken to the authors and they mention: "The main idea behind training the embedding network is indeed to adapt the representation to the task. A better way to think about it is, that you are trying to learn the metric for the nearest neighbor lookup, it's not exact though, as we don't reencode all the previous keys every time we update the network. The relationship between the metric for KNN lookup and our embeddings is really similar to what's happening with kernel trick with SVM's, the difference being that here we adapt the kernel." which suggests they don't update the keys.

[toth-adam]

Thank you very much for your exhaustive answer, I will toy around with the NEC agent implementation accordingly.

Best regards,

Adam

Based on the authors' comments, it seems they do update the keys by training the embedding network, which leads to a similar situation to using adaptive kernels. Although this metric is approximate since only some of the keys are updated. Also the name of the table as "differentiable" neural dictionary suggests that they back-propagate through the table. That's my understanding. Let me know how you think.

It makes me wonder if this framework will work in case we don't have an embedding network on top of observations, for instance in some gym experiments when we work directly with low-dim state space representation.

[EndingCredits]

My understanding is that it's a bit like an adaptive kernel , but where stored values aren't updated.

As for backpropogating though the table I don't really see how it can be done; I think it's just a poor choice of name, and what they really mean is they choose a differentiable distance metric (in MFEC they just take the average of the stored points, which is locally 'flat' so you can't get meaningful gradients). I could be wrong though.

What I might do is ask the authors if they could look over the code and see if it matches what they did, since they can't release their own code due to using some private internal libraries.

That would be great to ask, if you are in contact with the authors. Could you reproduce their results on some of the games they mentioned in the paper?

[toth-adam]

Actually, I've already contacted Alexander Pritzel (one of the original authors), here is his answer regarding this subject:

"The keys are updated through backpropagation, also when we reencounter an observation we overwrite the corresponding key."

The second half of the sentence corresponds to the case when the agent visits the exact same state again (they decide it by comparing the hash values of these states) in the forward pass; when this happens the state embedding key is overwritten in the dictionary (it may be slightly different, because the conv. network has been updated since).

Regarding implementation, they used a custom C++ OP in TensorFlow for approx. knn. (And for the gradient application on the DND keys as well, I guess.)

Regards,

Adam

[EndingCredits]

Hi,

Sorry for the late response.

Revisiting this issue it seems that my implementation is incorrect, and that the stored keys should be updated in the dict via backprop. However, it seems like updating the keys would require building the dictionary each time. I don't know yet if there is an easy way around this.

Best,
Will