Summary
The paper applies transformers to learning models for output sequences of several types of HMMs and Linear Dynamical Systems models.
The authors used two main tasks for evaluation -- next observation prediction and belief state prediction from the observations.
They compared the convergence speed of the transformer models with recurrent models and found, surprisingly (at least to me)
that RNN's converged faster and better than transformer models that they implemented.
The authors explored two main factors in the learning efficiency. Firstly they explored the effect of longer mixing length
on the ability of the transformer models to learn good models. Secondly they explored the effect of uncertainty in predicting
the belief state of the model from the observations, on the predictive capacity of the models. They found that when
the mixing length increased transformer models performed worse. They also found that when belief states were more strongly
correlated with next step predictions, then predicting belief states in the intermediate blocks could result in
improved performance by a block COT type approach.
Finally, the authors show some theoretical proofs related to -
1. The ability of RNNs and transformers to model sequences from HMMs with different levels of predictions based on depth
and length of sequences.
2. Existence of transformer models that can learn to model outputs from HMMs with arbitrary accuracy based on minimum
depth -- they used prior work which showed that transformers can model any arbitrary automata (of which HMM is but a sub-case).
Strengths
I thought the idea of using sequences from HMMs to study the learning abilities of transformers quite a good one -- it is possible to control the complexity of the sequences by the Markov-order of the HMMs; it is possible to control the number of states, and output distribution etc. The authors created Cyclic-{DET, RND,HARD} models which have different properties in terms of complexity, mixing length etc.
The block COT method seems to be good way to propagate more information through the model. During training this makes the model more autoregressive, which making it less computationally demanding than a fully next step autoregressive model based of the state of the model.
Weaknesses
I don't want to be too certain about what the weakness of the paper are, until I read the authors' rebuttals and clarifications, since its possible I may have missed some details and motivations.
My first reaction, is, sadly an instinctive one. In as much as it makes sense to me that deeper models should perform better, and that the more complex the sequence type, the harder it would be for an algorithm to model it. However, I am quite surprised, to see that recurrent models learn so much faster, and I want to be certain that this observation, which is repeated in various forms in the paper, holds. I have left several questions in the Questions section to help me understand the details. What I would have to say here is that it seems to me that transformers can mimic recurrent models, by learning to put all the attention in the previous step and previous step alone. In doing so, it can cause the models to flow information almost in a recurrent way. So I am puzzled that the recurrent model can outperform it.
Another thing that comes to mind is that I think to some extent the model studies not just the learning properties of transformers, but the combination of the transformers and the loss function. HMM's can often produce complicated, multi-modal output distributions. Modeling these with an L2-loss is likely to lead to over smoothed models. Such problems are well studied in ML, back to early days of using mixture density networks, where the model's output was a mixture model, and not just a single mean value. I suppose this is true also for the Recurrent models, but have the authors explored the impact of this on the learning ? I left some clarification questions below in the questions section to tease this apart a little bit more.
It is not clear why curriculum is important here. If I understand correctly, the same HMM is used over training a model. In that case, the complexity of each task is really just dependent on the property of the HMM, and the complexity is the same through the sequence. Is it not true that a shorter sequence from the HMM should be no more complex than a longer sequence from the same HMM ?
Experimental details are somewhat sparse; specifically what were the sizes of state space, n, what were the output dimensions ? Were continuous data normalized before training transformers. For discrete outputs, what were the output space (m) ?
Questions
I am afraid I had a lot of questions about the details that I think can be clarified, which will help me perform a better evaluation of the paper. I list them here in no particular order.
1. With your RNN's -- did you end up using ReLU RNN's ? It's surprising they perform so much better than transformers, even though LSTM's were typically the best in class RNNs years past.
2. I'm not entirely sure why you created the exact forms of HMM you did in this paper. The essence of the method seems to be to generate HMMs' that have different Markov-orders and different correlations between belief states and output sequences. Could this have been explored with other alternatives such as the complexity of the Markov Chains of the hidden sequences and the complexity of the emission distributions ?
3. Can you please provide more information about the experiments ? I am not sure what was the state sizes for different experiments,
4. With discrete state HMMs, and continuous outputs it would make sense that the output distribution would be a mixture model, rather than a single model. As such using an l2 loss may cause problems. Did the authors attempt to use mixture density network type approaches for these cases ?
5. Are the loss functions averaged over batch ? i.e. is loss divided by the batch size and sequence length ? I see only sequence length, in (B.1) and so I want to be sure.
6. B.1 lays you the loss function for training -- L2 for Matmul, RanLDS and cross entropy for other tasks. But evaluation on line 877 seems to suggest that evaluation loss is the l2 loss.. Is this also the case for categorical outputs ? I would think that the right evaluation objective would be the KL from the predictive distribution or the posterior belief state of the HMM itself. What is the case here ?
7. Do you have a study on the effect of the block size, b, of the block COT on the results ?
8. Epsilon on like 410 is overloaded with epsilon in line 344.
9. Line 455 - .."we can feed the output back into the transformer every b steps". Is this the actual output from the model ? If so, is gradient back propagated through this, and with what method (straight through, or reinforce, etc) ? If not, isn't it just teacher force model durning training, just like language models, and so training speed should be unaffected ?
10. During curriculum training, as I understand it here, you use increasingly longer sequences. Is batch size kept the same ? Are learning rate and other things adjusted for the changing batch size ? Does the training proceed from the previous point, or is learning rate warmup performed once more ?
11. In figure 6, why does the Cyclc-DET w.o. curriculum fall so significantly when going from 5 to 6 layers ?
12. It is really detailed to follow the HMM construction -- will you be releasing the code to create Cyclic-DET and Cyclic-RND ?
13. Any thoughts on why transformers perform almost a million times worse than a single layer RNN, for Matmult, Cyclic-DET and RND -- are these are with continuous outputs ? These results are quite surprising to me.