Unveiling Markov heads in Pretrained Language Models for Offline Reinforcement Learning

Thank you for your detailed review and feedback. We appreciate your positive comments about the novelty and presentation of our work. Please kindly find the response to your concerns below.

W1. For figure 1 and figure 4, DT is tested on only one task.
We have tested DT and GPT-DT on more tasks and the results are shown in Fig.1~10 at [URL]. The conclusion are consistent among different tasks. We will include these results from multiple tasks in the final version.

W2. Stronger results for Theorem 4.3
We extend our proof to show high probability bound. By Theorem 4.3 (page 4), we know $\frac{\mathbb{E}[D]}{\mathbb{E}[O]} > r$ for some $r>0$, where $D \triangleq \frac{1}{K} \sum_{i=1}^K |\Pi_{ii}|, O \triangleq \frac{1}{K(K-1)} \sum_{i \ne j} |\Pi_{ij}|, {\rm and\ } \Pi\triangleq E A E^T.$ We assume each element in embedding vector is i.i.d sampled from $\mathcal{N}(0,1)$. Then $|\Pi_{ij}|$ is sub-exponential with sub-Gaussian proxy $||A||_F$. By Bernstein inequality and Hanson-Wright inequality, we can get the bound for $\mathbb{P}\left( |D - \mathbb{E}[D]| \geq \varepsilon \right)$ and $\mathbb{P}\left( |O - \mathbb{E}[O]| \geq \varepsilon \right)$. By choosing $\varepsilon\triangleq \frac{\mathbb{E}[D] - r \mathbb{E}[O]}{2r} > 0$, we have $\frac{D}{O} > \frac{\mathbb{E}[D] - \varepsilon}{\mathbb{E}[O] + \varepsilon} > r$. Due to $\mathbb{P}(\frac{D}{O}\leq r) \leq \mathbb{P}(D < \mathbb{E}[D]-\varepsilon) + \mathbb{P}(O> \mathbb{E}[O] + \varepsilon) \leq \delta_0$, we can obtain the high-probability bound theorem. We will add the detailed theorem and proof in the final version.

W3. Explanations for the scale$K$
Recall that in Theorem 4.5 (page 5), for any K<\Big\lfloor\min\Big\\{\frac{\rho}{r+1}\frac{\overline{|A_{ij}^0|}}{\eta_0 B},\frac{A_{ii}^0}{\eta_0 B}{\rm for\ }i=1,\cdots,d\Big\\}\Big\rfloor, Markov heads preserve. In remark 4.8, during the observations in experiments, we found that $\eta_0 = 1e-4$, $B = 1e-6$, $\overline{|A^0_{ij}|} = 1e-2$ and we set $r=20$ then $\rho = 3.61$. Since $A^0_{ii}$ are larger than $\overline{|A^0_{ij}|}$, we can ignore the second part in this inequality. By all the parameters, we can conclude that for any $K < 1.7 e7$, theorem 4.5 will be satisfied. In our setting, we let $K = 1e4$ which is within the theoretical range.

Q1. In long-term task, why GPT-DTMA cannot surpass DT?
In our theoretical results, we demonstrate that Markov heads tend to disproportionately attend to the final input token, thereby impairing long-term planning capabilities. Although the incorporation of a gating mechanism can attenuate the influence of these heads, it does not eliminate them entirely. Another contributing factor may be the difference in training corpora: DT is trained exclusively on a reinforcement learning corpus, whereas GPT-DTMA fine-tunes weights pre-trained on a language modeling corpus. Furthermore, the number of trainable parameters plays a nontrivial role in model performance. Consequently, DT tends to exhibit superior empirical performance. The primary motivation for proposing GPT-DTMA is to enhance the long-term planning ability of GPT-based decision transformers while maintaining lower computational costs by circumventing training from scratch.

Q2. Would the positional embedding play a role in yielding Markov heads?
We acknowledge that positional embeddings may also contribute to the model’s tendency to attend more to the final input token, however, positional embeddings are not directly related with Markov heads. It is important to note that neither DT nor GPT-DT employs positional embeddings. Our analysis is confined to the matrix product $W_q W_k^T$, which is independent of any positional encoding. We define $W_q W_k^T$ as the Markov matrix. In Theorem 4.3 (page 4), we demonstrate that for any random embedding matrix $E$, the transformed matrix $E W_q W_k^T E^T$ remains a Markov head in expectation.

Thanks again for your constructive suggestions and we will include them in the final version. We are looking forward to further discussions.

Thank you for your comprehensive review and constructive comments. We appreciate your acknowledgement regarding the originality, significance, clarity and presentation of our work. Please find our response to suggestions and questions.

S1&Q3. Test other PLMs to confirm generality.
We have tested the existence of Markov heads in other pre-trained large models, such as GPT-J, ImageGPT [1]. We examined the attention heads in GPT-J and found that all of them cannot satisfy the condition (i) in Definition 4.1, i.e. Markov head is not detected. We also test the performance of initializing DT with GPT-J and ImageGPT checkpoints for short-term environments. Table R1 shows the result comparisons among different PLMs. Results show that without Markov heads, the performances of GPTJ-DT and ImageGPT-DT fail to align with GPT-DT.

Table R1:

[1] Chen, Mark, et al. "Generative pretraining from pixels." International conference on machine learning. PMLR, 2020.

S2. Analyze MoA behavior per time step (how weights change dynamically).
We analyzed the step-wise changes of average weights for Markov heads and non-Markov heads. Fig.11 and Fig.12 in [URL] show the results for short-term and long-term environment. We conclude that, for short-term environment, the weights of Markov heads gradually ascent while the weights of non-Markov heads gradually descent. In contrast, for long-term environment, the weights of Markov heads gradually descent while the weights of non-Markov heads ascent.

S3. Compare GPT-DTMA’s performance against value-based offline RL baselines.
We have conducted experiments on value-based offline RL methods (CQL) and Decision Convformer (DC). Table R2 shows results for short-term environments, in which CQL generally perform worse than GPT-DT or GPT-DTMA, and DC outperforms GPT-DTMA. Table R3 shows results for long-term environments, where CQL outperforms other methods and DC performs worst. While the results of DC imply that directly utilizing the Markov property benefits short-term environment performances and hurts long-term performances, they strengthened our claim.

Table R2:

Table R3:

S4. Extra related works
Thanks for pointing out. Decision Convformer (DC) explicitly trains convolutions to enhance local attention; Trajectory Transformer employs a different trajectory formation than DT. In contrast with these work, we identify Markov heads in PLM that influences the short-term and long-term planning ability. We will include these discussions in our final version.

Q1. How was the threshold for Markov heads chosen?
We begin by recalling the definition of a Markov matrix as stated in Definition 4.1: all diagonal elements are positive, and the ratio $\frac{\overline{|A_{ii}|}}{\overline{|A_{ij}|}} > r$. Let $\overline{|A_{ii}|} = m$ and $\overline{|A_{ij}|} = n$, then after applying softmax function, the diagonal element becomes $\frac{e^m}{e^m + (d-1) e^n}$, where $d$ is the embedding dimension. To determine a reasonable range for $r$, we assume that the diagonal element $\frac{e^m}{e^m + (d-1) e^n}$ should be at least $0.5$ in order to express higher attention on last input token, i.e., $\frac{e^m}{e^m + (d-1) e^n} > 0.5$. Then we can obtain that $r > \ln(d-1) + 1$. In our setting, $d = 768$ and $r > 7.64$. In remark 4.7, we set $r=20$ to identify Markov heads of GPT-DT. Under this threshold, the corresponding diagonal element is at least 0.99, indicating an extreme focus on last input token.

Q2. Would a manual reduction of Markov head influence improve long-term tasks further?
A manual reduction of Markov head influence is shown to improve performance on long-term tasks. We investigated this approach using Attention with Reverse Linear Biases (ALiBi-R). ALiBi-R introduces a manual reduction of Markov head behavior by penalizing diagonal elements and simultaneously enhancing the weights assigned to distant elements in the attention matrix.

Table R4:

For the minor clarity issues, we will revise and address this issues accordingly.

Thank you again for your constructive comments!

Thanks for your valuable comments. Please kindly find the response to your concerns below.

Q1. Would a MoE (mixture of experts) perform better in these long-term environments?
It’s interesting to investigate whether MoE perform well in long-term environments, however, MoE is not considered in this work due to the main focus of this work is on unveiling the impact of Markov heads in PLMs that affect offline RL performances. Indeed, Markov heads may exist in each expert of an MoE model, so our findings may also provide insights for explaining the performance of MoE in long-term environments.

Q2. What is the high-level intuition behind a Markov head? Additionally, why is diagonal dominance important in the context of NLP and RL?
The underlying intuition behind Markov heads is that they function as strong local policy learners, relying predominantly on the most recent observation to inform decisions. This behavior mirrors the memoryless property of Markov processes. Diagonal dominance is a critical characteristic, as it ensures that attention is primarily focused on the current input token, thereby reinforcing the Markovian assumption and promoting stable learning in highly reactive environments. While we acknowledge that not all attention heads exhibit diagonal dominance, we demonstrate that this variation reflects the emergence of diverse temporal planning behaviors. Specifically, only a subset of heads learns to follow short-term decision patterns, a phenomenon particularly pertinent in reinforcement learning settings where tasks may exhibit heterogeneous temporal dependencies.

Q3. How did you use for the train and test splits when training DT and GPT-DT? Did you train and test on different datasets from the same environment, or did you train in one environment and test in another?
Following common practice in offline RL research, we train each model with offline datasets collected by D4RL for each environment respectively, and test the model in the same environment.

Q4. How do the non-Markov heads affect performance in both long-term and short-term environments? In general, what percentage of heads are considered Markov heads? Are Markov heads fixed within a particular environment, or do the heads classified as Markov change depending on the environment?
Non-Markov heads show almost equal attention on each tokens, and we can see that the weights of all non-Markov heads become larger in long-term environments, such as PointMaze tasks. During our observations, we found that 3 of 12 heads are Markov heads. The Markov Heads are defined according to Definition 4.1 and 4.2 regardless of environments. According to Theorem 4.5, for GPT-DTs, this Markov Head property are fixed and won’t change by fine-tuning. For DT, the model is trained to obtain or not obtain Markov Heads based whether it’s short-term or long-term environment.

Q5. Is the issue discussed in Section 4.2 stem from distribution shift problems due to changes in the training and testing environments or related to modeling flaws.
We would like to clarify that the issues discussed are neither a distribution shift problem nor related to modeling flaws. We state that when given an new task for a PLM initialized DT model to be fine-tuned on, we may not know the planning ability required (either short-term or long-term) for this task in advance. If long-term planning ability is required, we have proved that the Markov Heads in PLM weights will hurt performance.

Q6. Why do you need to adaptively control the weight of the attention heads?
Due to the issue stated in Section 4.2 and given our statement in Q5, we would like the model to automatically adapt to the planning ability needed. Therefore, GPT-DTMA learns a weight reducing the influence of Markov heads to enhance long-term planning ability of GPT-DTs.

Q7. Do you have a skyline model that shows what good behavior looks like?
While our experiments follow the standard offline RL settings, trainings and testings are conducted on the same task for each model. Therefore, the current results have represented the “good behavior” given each model configuration.

Q8. How does the Markov heads change for DT and GPT-DT as in Table 5/6?
We would like to clarify that Markov head is not observed in DT. Table 5/6 shows the head weights given by the MoA module in GPT-DTMA, therefore they are not available for DT and GPT-DT.

Q9. With an increase in context length would the models tend to focus more easily on distant information?
In Theorem 4.3 and 4.5, we prove that Markov heads persists in GPT-DT during fine-tuning, so it will not tend to focus on distant information natively, even when context length is increased. Therefore, the performance gain brought by MoA to GPT-DTMA is non-trivial and cannot be replaced by simply extending the context length.

Thank you again for your reviewing efforts and we sincerely hope our responses have addressed your concerns.

Thanks for your valuable comments. Please kindly find the response to your concerns below.

W1. The improvement of GPT-DTMA is within the confidence interval.
The experiment are repeated three times to ensure significance. We would also like to emphasize that, one important objective for our experiments is to validate the influence of Markov heads in short-term and long-term environments. Since Markov heads play important role in both GPT-DT and GPT-DTMA, it is reasonable that their performance stay close in short-term environments. However, their performance gap in long-term environment is larger, showing support to our main claim.

W2. The paper doesn't extensively compare its absolute performance against the broader state-of-the-art in offline RL methods on these benchmarks.
We have extended experiments to compare with value-based offline RL methods (CQL) and Decision Convformer in our results. Please refer to Table R2 and Table R3. Table R2 shows results for short-term environments, in which CQL generally perform worse than GPT-DT and GPT-DTMA and DC outperforms GPT-DTMA. Table R3 shows results for long-term environments, where CQL outperforms other methods and DC performs worst. While the results of DC imply that directly utilizing the Markov property benefits short-term environment performances and hurts long-term performances, they strengthened our claim.

Table R2:

Table R3:

Q1. Your analysis focuses primarily on GPT-2. Have you conducted preliminary investigations or do you have hypotheses about whether similar Markov head phenomena exist and exhibit the same stability when using other PLM architectures?
We have tested the existence of Markov heads in other pre-trained large models, such as GPT-J, ImageGPT [1]. We examined the attention heads in GPT-J and found that all of them cannot satisfy the condition (i) in Definition 4.1, i.e. Markov head is not detected. We also test the performance of initializing DT with GPT-J and ImageGPT checkpoints for short-term environments. Table R1 shows the result comparisons among different PLMs. Results show that without Markov heads, the performances of GPTJ-DT and ImageGPT-DT fail to align with GPT-DT.

[1] Chen, Mark, et al. "Generative pretraining from pixels." International conference on machine learning. PMLR, 2020.

Table R1:

Q2. The Markov heads are based on a ratio r (in Definition 4.1/4.2 context). How sensitive are your findings – specifically, the number of heads identified as Markov and the overall conclusions about their impact – to the choice of this threshold r?
We begin by recalling the definition of a Markov matrix as stated in Definition 4.1: all diagonal elements are positive, and the ratio $\frac{\overline{|A_{ii}|}}{\overline{|A_{ij}|}} > r$. Let $\overline{|A_{ii}|} = m$ and $\overline{|A_{ij}|} = n$, then after applying softmax function, the diagonal element becomes $\frac{e^m}{e^m + (d-1) e^n}$, where $d$ is the embedding dimension. To determine a reasonable range for $r$, we assume that the diagonal element $\frac{e^m}{e^m + (d-1) e^n}$ should be at least $0.5$ in order to express higher attention on last input token, i.e., $\frac{e^m}{e^m + (d-1) e^n} > 0.5$. Then we can obtain that $r > \ln(d-1) + 1$. In our setting, $d = 768$ and $r > 7.64$. In remark 4.7, we set $r=20$ to identify Markov heads of GPT-DT. Under this threshold, the corresponding diagonal element is at least 0.99, indicating an extreme focus on last input token. From Table 1, we can see that there are three Markov heads before/after fine-tuning for any $r\in(20,+\infty)$.

Thank again for your thoughtful reviews. We will incorporate the discussions in the revised version and we are looking forward to further discussions.