Unveiling Markov heads in Pretrained Language Models for Offline Reinforcement Learning
Authors: Wenhao Zhao
Qiushui Xu
Linjie Xu
Lei Song
Jinyu Wang
Chunlai Zhou
Jiang Bian
Venue: ICML 2025 poster (2025)
Keywords: decision transformer; cross-domain
Primary area: reinforcement_learning->planning

Abstract: Recently, incorporating knowledge from pretrained language models (PLMs) into decision transformers (DTs) has generated significant attention in offline reinforcement learning (RL). These PLMs perform well in RL tasks, raising an intriguing question: what kind of knowledge from PLMs has been transferred to RL to achieve such good results? This work first dives into this problem by analyzing each head quantitatively and points out Markov head, a crucial component that exists in the attention heads of PLMs. It leads to extreme attention on the last-input token and performs well only in short-term environments. Furthermore, we prove that this extreme attention cannot be changed by re-training embedding layer or fine-tuning. Inspired by our analysis, we propose a general method GPT-DTMA, which equips a pretrained DT with Mixture of Attention (MoA), to enable adaptive learning and accommodate diverse attention requirements during fine-tuning. Extensive experiments demonstrate the effectiveness of GPT-DTMA: it achieves superior performance in short-term environments compared to baselines, significantly reduces the performance gap of PLMs in long-term scenarios, and the experimental results also validate our theorems.

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
Wenhao Zhao
1*t Qiushui 2
Linjie Xu 3 Lei Song"
Jinyu Wang* Chunlai Zhou'
Jiang Bian 4
Abstract
2018), computer vision (Yuan et al., 2021) and AI4Science (Wang et al., 2024) tasks, for it encodes the input data into powerful features via the attention mechanism. Applying transformers to the field of reinforcement learning (RL), De- cision Transformer (DT) (Chen et al., 2021) which models the offline RL problem (Levine al., 2020) to return- conditional sequence-to-sequenne problem, has demon- strated superior performance in solving different types of RL problems (Correia & Alexandre, 2023; Mezghani et al., 2023; Badrinath et al., 2024; Janner et al., 2021).
Recently, incorporating knowledge from pre-
trained language models (PLMs) into decision transformers (DTs) has generated significant at- tention in offline reinforcement learning (RL). These PLMs perform well in RL tasks, raising intriguing question: what kind of knowledge from PLMs has been transferred to RL to achieve such good results? This work first dives into this problem by analyzing each head quantitatively and points out Markov head, a crucial component that exists the attention heads of PLMs. leads to extreme attention on the last-input token and performs well only in short-term environments. Furthermore, we prove that this extreme atten- tion cannot be changed by re-training embedding layer or fine-tuning, Inspired by our analysis, we propose a general method GPT2-DIMA, which equips a pretrained DT with Mixture of Attention (MoA), accommodate diverse attention require- ments during fine-tuning. Extensive experiments corroborate theorems and demonstrate the ef- fectiveness of GPT2-DTMA: it achieves compara- ble performance in short-term environments while significantly narrowing the performance gap in long-term environments.
Recently, using pretrained language models (PLMs) and fine-tuning them for specific tasks (Kenton & Toutanova, 2019; Raffel al.. 2020) has gradually replaced training Transformers from scratch. Pretraining phase enables PLMs to gain rich universal language representations transferable to various downstream tasks (Xiong et al., 2024). Remark- ably, Noorbakhsh et al. (2021); Goel et al. (2022) showed that PLMs can effectively tackle tasks beyond NLP. Inspired by the powerful knowledge transfer capabilities of PI.Ms, Reid et al. (2022); Shi et (2024) attempt to transfer NLP knowledge RL domain, indicating that PLMs can en- hance the performance of downstream offline RL tasks as well. We collectively refer to the DT variant models, which are initialized with PLMs, as GPT2-DT.
To understand this surprising efficacy when transferring knowledge from PLMs to offline RL, Shi et al. (2024) con- cluded that sequential modeling ability is the key to un- leashing the potential of PL.Ms for RL. tasks. Furthermore, Takagi (2022) proposes hypothesis that pretraining with text is likely to make PLMs get context-like like information and utilize it to solve the downstream task. However, it is still unclear that why these key features are essential for RL tasks and how to acquire these knowledge through pre- training. To bring clarity to PLMs in offline RL and provide more insightful guidance how better train and adapt transformer-batt models for RL tasks, we dive into the mechanism of PLMs and do more quantitative analysis on the direct transferred parameters to RL. Based on our quan- titative analysis, we first propose the definition of Markov heads, which inherit from PLMs and are fundamental for RL tasks and then adapt PLSs to all kinds of RI. tasks by flexibly utilizing the properties of Markov heads.
1. Introduction
Transformers (Vaswani et al., 2017) achieve significant im- provements in natural language processing (Devlin et al.,
'Equal contribution; the order is decided by dice roll. 'This work was conducted when the author
intern at Microsoft Research Asia CCmmpueer Science Department, Renmin Uni- versity China, Beijing, China 'Department Industrial and Manufacturing Engineering, Pennsylvania State University, Uni- versity Park. PSA "Scoool of Electrical Engineering and Com- puter Scientce, Queen Mary University of London, London, UK "Microsoft Research Asia, Beijing, China. Correspondence to: Lei Song <lei.song @ microsoft.com>..
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
DT
GPT2-DT
Transformer Attention Layer
Transformer Attention Layer
Similar Attention Score Distribution
Fower licrations
More lnenttions
Fewer Iterations
More Iterations
013
Figure 1. We compare the attention score distribution of DT and GPT2-D5 during the process of fine-tuning within fewer iterations and more iterations under Hopper-medium environment, More comparisons of the attention score distribution on Walker2d-medium and Halfcheetah-mmdiim can be found in Appendix B.
In this work, we investigate both short-term and long-term planning ability (Shang et al., 2022) of PLMs. To be more specific, we define short-term environments as where the actions prediction recalls a small amount of contextual in- formation from most recent timesteps (e.g. MuJoCo Loco- motion), and long-term environments that require the agent to extract useful information from the whole previous con- text (e.g. PointMaze), as context from far-off timesteps is helpful to predict reasonable actions under this scenario. The reason why we want to consider both short-term and long-term is that, our investigation found while PLMs per- form better in short-term environments, it is less effective in long-term environments compared to a DT trained from scratch. Therefore, we are interested in what kind of NLP knowledge from PLMs has been successfully transferred to short-term RL tasks but failed in long-term RL task.
propose GPT2 combining PLMs with Mixture of Attention (MoA). Here, we consider each attention head as expert and learn the weights of each expert adaptively ac- cording to each downstream RL task, so that GPT2 -DTMA can automatically identify the required planning abilities by adjusting the weights of Markov heads when predicting the next action. CFT2-DTMA achieves comparable perfor- mance in short-term environments and narrows performance gap long-term environments.
In summary, our contributions follows:
We first propose the definition of Markov heads and fur- ther explore the properties of Markov heads. We prove that Markov heads will be preserved under any embed- ding initializations and fail to modify by fine-tuning. We also validate our theorems through extensive com- parative experiments.
We discover that, as shown in Figure 1, GPT2-DT (Reid et al., 2022; Shi et al., 2024; Yang & Xu, 2024) has shared similar attention score distribution with well-trained DT, even at the initial stage of its training. We define the im- portance score of each head by zero-out each head and measuring the difference of the predictions, We also demon- strated the heatmap of each head and found that all attention heads can be divided into two categories. In particular, we define Markov heads as one kind of special heads inheriting from NLP tasks and identify that these heads lead DTss to focus more on the most recent state in short-term RL tasks by our theoretical analysis. Furthermore, we examine that Markov heads will be retained under any random embedding initializations and cannot be easily switched to Non-Markov heads by fine-tuning both theoretically and practically.
We do a comprehensive quantitative analysis on the heads of PLMs by using different ways to explore the deeper mechanism of PLMs based on our definitions and theoretical analysis. We show that Markov heads are the key information transferred from NLP.
Finally, we propose general model GPT2 DTMA that can learn the weights of each head automatically and it significantly narrowed the performance gap of our model between short-term and long-term RL tasks com- pared to baselines.
2. Related Work
Naturally, we are also curious if we can still adapt PLMs to all kinds of RL tasks without training from scratch even though Markov heads will not be changed by fine-tuning. To bridge the performance gap in long-term environments, we
2.1. Decision Transformer
Transformer-based architecture (Vaswani et al., 2017) is widely explored in offline RL. Decision Transformer (DT) (Chen et al., 2021) firstly models RI. as an autoregressive
2

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
generation problem. By conditioning on previous trajec- tories and returns-to-go, DT can generate desired actions without explicit reward modeling or dynamic programming. Furthermore, Furuta et al. (2021) explores other kinds of hindsight information instead of returns-to-go that can ben- efit sequential cerision-making Q-learning DT (Yamagata et al., 2023) proposes to combine DT with dynamie program- ming by utilizing a conservative value function to relabel returns-to-go in the training dataset. In addition modi- fying the context by augmented information (Wang et al., 2024), some other work is also attempting to modify the model architecture of DTs. Shang et al. (2022) argues that the DT structure, which requires all tokens as input, is inef- ficient for learning Markovian- like dependencies. Kim et al. (2023) proposes replacing the attention layers with convolu- tional layers better capture the inherent local dependent patterns in trajectories of RL tasks. However, all existing work focuses mainly on data augmentation for whole of- fline datasets and model redesign. These methods are still and unstable
information to assist DT' in action prediction.
3. Preliminaries
3.1. Offline Reinforcement Learning
In offline RL, the policy model π is trained on a pre- collected dataset, rather than interacting with the environ- ment in real time. We consider learning in a Markov deci- sion process (MDP) described by M = (po. (Po.S,A, A, P. R, 7), where po is the initial state distribution, S is the state space, A is the action space, P(s'| ss a) is the transition probability, R(s, a) is the reward function and 7 E (0, 1] is the discount factor. We use St, ag and re = R(st, st. a.) to denote the state, action and reward at timestep t, respectively. The goal of offline RL is to find an optimal policy x that maximizes the y-discounted expected return:
max ESs PPP
Σ t=
2 R' ai).
(I)
struggling with low computation-efficienn performance on all RL tasks. In our work, we aim at with- out modifying the model architecture or introducing new hidden information, and prove that PLMs can effectively capture local information thanks to Markov heads.
3.2. Decision Transformer
DT's (Chen al., 2021) are designed to generate next ac- tions relied on the information about expected future re- wards for the current trajectories. therefore they propose returns-to-pe go (RTG) denoted as RS - Ei= Ti, which rep- resents the sum of expected future rewards from the current timestep t. In order to incorporate RL problems into a sequential model, DTs transform the trajectory into the for- mat of (Ro, So, ag. R1, $1, 01...... Rr, ST, aT), where Re, 8t, a the RTG, state and action at timestep f, respec- tively. The next action predicted by policy To is conditioned on the previous trajectory T up to the timestep t 1, current RTG Rg and current state 8y:
2.2. Pretrained Methods for Offline RL
A variety of approaches have been proposed to leverage pretraining technique in offline reinforcement learning. Nair et al. (2020); Zheng et al. (2022); Kostrikov et al. (2021) follow the standard Seetrain-fintune paradigm, pretrain- ing on a single-task offline dataset followed by fine-tuning in an online environment. Xie et al. (2023) pretrain deci- sion transformers (DTs) using reward-free trajectories, and then fine-tune the model using reward-tagged trajectories collected online. Additionally, pretraining on mixed data collected from various RL tasks (Liu et al., 2022; Xu et al., 2022; Sun et al.. 2022; Carroll et al., 2022: Wu et al., 2023) can achieve good performance on downstream tasks through fine-tuning or few-shot learning.
d; = ne (Ro, so, ao, Re, st), and Ta is trained by minimizing the squared error between the true action de and the predicted next action aA:
T
CDT
Σ t=
||:
at
(3)
However, the scale of pretraining datasets with trajectory structures in RL domain is not comparable to the scale of text-structured datasets in NLP domain, which limits DT's ability to thoroughly explore the nature of RL tasks. To ad- dress this, a pretraining paradigm on language datasets has been proposed; more details can be found in Appendix A. Reid et al. (2022); Shi et al. (2024) suggest that initializ- ing DT's with pretrained parameters from PLMs, such as those from GPT-2, can overcome this limitation and outper- form DT's that are initialized randomly. Yang & Xu (2024); Zheng et al. (2024) also utilize text content pretraining samples, but during the fine-tuning phase, they incorporate the concept of Prompt-DT (Xu et al., 2022) by using tra- jectories from the single-task or multi-tasks as prior input
4. Methodology
In this section, we will first give the definition of Markov heads and analyze the properties of Markov heads. Based on our theoretical analysis, we then introduce a general approach
that shows great adaptive abilities to either long-term or short-term environments.
4.1. Markoy Head
Definition 4.1. For any matrix A E gdxd matrix if and only if A satisfies the following conditions:
df A is a Markov
3

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
The diagonal elements {A,,)i= are all positive.
(i). The ratio between the mean of absolute diagonal ele- ments and the mean of absolute off-diagonal elements are strictly greater than some large positive number
r, ic.,
Ag
where |Ai]
-d And and
6.24
Aul
AGE
Aij
0.00
For each head, have three weight matrices corresponding to key, query and value. According to the definition of atten- tion scores, we will focus our analysis on W' which is the matrix multiplication between key weight matrix and query weight matrix. Then, we introduce the definition of Markov head based on Definition 4.1.
14
Figure 2. Importance score of cach head of GPT2-DT in Hopper- medium environment. The red bar represents the importance scores allocated to the Markay head in the first attention layer. We also demonstrate importance score in PointMaze-large in Appendix D.
Definition 4.2, For head i, if the weight matrix W (T a Markov matrix, then we say head i is a Markov head.
Proof. See more details in Appendix E.2.
4.1.1. THEORETICAL ANALYSIS
In the next theorem, we can show that Markov matrix will be maintained after finite-time fine-tuning.
Markov matrix has some important properties that can give us more insight about what knowledge in PLMs is benefi- cial for RL tasks or limit the performance of generalization ability in RL tasks. In this section, we provide a theoreti- cal analysis and detailed explanation of the properties of Markov matrix and Markov heads.
Theorem 4.5. Suppose all the norm of the gradients are upper bounded by some positive number B and the learning rate Tis has a uniformly upper bound no- For any Markov matrix A for any training iterations K <
[min {_ 14%1 A for 1, .d]], AK is still r+I B B
Theorem 4.3. For any random embedding vector e, E RIxd = 1,.... K, the elements of each are i.i.d. sam- pled a normal distribution. For any given Markov ma-
Markov matrix.
Proof. See more details in Appendix E.3
trix A E Rdxd satisfying Definition 4.1, then E[EAE"] is also a Markov matrix, where E- = (01, ""ek""
,Ck)² E RRxdd
Remark 4.6. According to Theorem 4.5, we can conclude
that during the stage of fine- fine-tuning on the specific down- stream RL tasks, Markov heads preserve and they can never be switched to Non-Markov heads.
Proof. See more details in Appendix E.I.
Based on Theorem 4.3, we extend the expectation result to a high probability bound. As corollary, we show that for any random embedding matrix E, any given Markov matrix A satisfies Definition 4.1 with high probability.
4.1.2. QUANTITATIVE ANALYSIS
In this section, we conduct quantitative analysis on the weight matrices transferred from PLMs to examine the cxis- tence of Markov heads in PLMs and their properties.
Corollary 4.4. Define the diagonal mean and off-diagonal mean II
II A EAFT as
First, we measure the importance of each attention head. We initialize DT with pretrained GPT2-small parameters (GPT2-DT) DT) and fine-tune it on RL tasks (Reid et al., 2022; Shi et al., 2024). Then, we test GPT2-DT under online environments, obtaining the predicted action à. To obtain the importance score each head, we replace the output of the head i with a zero vector (Bau et al., 2020; Michel et al., 2019) and notate the corresponding predicted action after removing head as å_₃. We define the importance score of the head as ||ā - 2=1// The results are visualized in in Figure 2.
D A K
Dub
o
Mayl-
(4)
K(K - 1)
ifj
Given a Markov matrix A E IRddd satisfying Definition 4.1, for any A -rD]-rB[( O such that 2r
0. there exists constants C1, C2
P =r) 21 (k)
(5)
where ē is defined as min( TA ) and || is the Frobenius norm of Markov matrix A, defined as || ET
It is notable that some heads exhibit higher importance scores than the others. To further explore the inner property of these heads, we focus our analysis on the weight matrix

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
Non-Markov
Markov head
Table 1. We analyze the weight matrix WP " of each head GPT2 without fine-tuning (Initialization) and after fine-tuning (Fine-tuned). Details of (i) and (ii) can be found in Definition 4.1.
0.73
0.50
0.0
Head Index
Initialization
Fine-tuned
0.25
(ii)
(i)
(ii)
-0.50
x
X
V
V
/
V
Figure 3. The heatmaps of W (W;)7" W. and Markov head (right).
in Non-Markov head (Icft)
X
X
X
3
V
X
V
X
X
X
fine-turing
fine-tuning
V
V
X
X
x
X
X
X
Non- Markow Head
X
8
X
X
X
10
V
11
X
X
X
(W of each head i. From the statistical data in Table and the heatmaps visualized in Figure 3, we can conclude that these heads with higher importance scores are exact Markov heads. Indeed, Markov heads exist in both PLMs and GPT2-DT, and such heads play a significant role next-action prediction tasks.
Morkon
Remark 4.7. Let m and n denote the mean of the absolute diagonal and dff-diagonal elements Definition 4.1. After applying the softmax function, the attention weight on the diagonal becomes (d)= where d the embedding eth
Figure 4. The two sub-figures above below score matrices for Non-Markov head Markov head before (left) and after fine-tuning (right). Lighter colors represent higher atten- tion scores, while darker colors represent lower scores.
attention
dimension. To determine a reasonable range for r, we as- sume (d== should be least 0.5 ensure that the head focuses more on the most recent input token. Solving this inequality yields > In(d 1) + In our setup, the embedding dimension d = 64 and set 8 to identify Markov heads in GPT2-DT.
random embedding matrix, the property of extreme attention to the last input token preserves in expectation. Figure 4 further confirms our theorems that the property of Markov Heads holds before and after fine-tume nee-tuning.
To verify both Theorem 4.3 and Theorem 4.5, we compare the attention score distributions before and after mne-tuning of both Markov heady and Non-Markov heads. Figure visualizes the comparison with heatmaps. On the one hand, observe that the attention distribution of Markov heads prior to fine-tuning is consistent with Theorem 4.3, Even when the embedding vectors for states, actions, and returns- to-go are randomly initialized, Markov heads exhibit strong attention to the last input token. This aligns with our theoret- ical result that EEee Ae?" >> Aee", where the last token corresponds the current state st. On the other hand, we find that after fine-tuning, the weight product W (WW)"" for all Markov heads in GPT2-DT continues to satisfy the definition of a Markov matrix, providing empirical support for Theorem 4.5.
Remark 4.9. During the fine-tuning process, we set our learning rate as le-4 and used 10K steps for warming up, so the upper bound пo approximately 1c-4. We observe that the magnitude of the maximum norm gradient of W' (W"kkr is the order of 1c-6, while the average absolute value of the off-diagonal elements is on the order of 1e-2. Given our setting of = 8, these observations indicate that the assumptions required by Theorem 4.5 are easy to satisfy under standard hyperparameter configurations.
4.2. GPT2-DTMA: A General Approach with Adaptive Planning Ability
Our analysis of Markov heads in GPT2-DT reveals that their extreme attention the last-input token aligns closely
Remark 4.8. By Theorem 4.3 we show that under any
5

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
MoA Ougutyy
Table 2. We compare and GPT2-D21 with mean normalized
score (1) in MuJoCo Locomotion and mean test episode length (U) in PointMaze respectively. Dataset names are abbreviated as follows: medium m, medium-replay as m-r.
Alterion leea Weight
Environment
Dataset
DT
GPT-DT
Gate Controller
Rrad,
Manic
denda
Hopper-m
67.4
77.9
MuJoCo Locomotion
Hopper-m-r
74.1
77.9
Walker2d-m
74.3
77.1
Inpute e
Walker2d-m-r
71.9
74.0
Figure 5. Mixture of Attention heads (MoA) allocates different weights to each attention head through gate controller network.
umaae
56.3
59.7
PointMaze
medium
158.7
188.0
large
195.3
257.3
with the behavior of memoryless policies (Littman, 1994).
can be expressed as:
However, such behavior may limit performance in long-term environments, where past trajectories carry rich contextual information that is crucial for accurate action prediction. In addition, it is often unclear whether long-ten term-teemmmm or short-term planning is required without prior knowledge of the environ- ment. To address this challenge, we extend GP GPT2- DT and propose general framework with adaptive planning abil- ity, GPT2-DTMA, to handle both short-term and long-term environments adaptively.
N u(e) Σ Gi(e) (e)
(6)
The dense gate controller network G is implemented as a linear layer W, followed by softmax activation:
G:(e) softmax (W)
(7)
From the results in Table we observe that GPT2-DT performs well in MuJoCo Locomotion (short-term) environ- ments, but exhibits significant performance drop in Point- Maze (long-term) environments, where it underperforms compared to the baseline DT. This suggests that the effec- tiveness of Markov heads is inherently biased toward envi- ronments with specific planning ability requirements. This raises natural question: How can improve GPT2-DT so that it retains its advantages in short-term scenarios while closing the performance gap in long-term environments?
where Ge(e) is the trainbale weight assigned to head and (e) denotes the output of attention head with query projection W keyword projection W and value projec-
tion WT. (ead,(e) is defined as:
(cW (EW!))
Head, (e) softmax(
(EW?), (8)
Vdl
where de denotes the dimension of the embedding layer.
5. Experiments
To address this key challenge, a natural idea is to adaptively control the influence of Markov heads based on the implicit requirement for planning ability each environment. While the feed-forward network (FFN) following the attention layers is capable of integrating information, our analysis in Appendix F shows that it fails to effectively balance the contributions of Markov and Non-Markov heads. To overcome this limitation, we propose enhancing GPT2-D' with a Mixture of Attention (MoA) mechanism (Zhang et al., 2022; Jin al., 2024), which results our method named GPT2-DTMA. An illustration of the proposed architecture shown in Figure 5. MoA treats each attention head as an expert specialized in capturing certain features temporal dependencies. A trainable gating network G; is introduced to adaptively assign weights to each head i's output. Given an input embedding vector c, MoA computes a weighted sum over all attention head outputs, where the weights are dynamically determined by G, Formally, the MoA output
From the previous analysis, we identified the existence of Markov head both PLMs and GPT2-DT and showed that they assign higher attention weights to the current state. However, it remains unclear whether this property bene- fits all reinforcement learning tasks. Prior studies (Reid et al., 2022; Shi et al., 2024; Yang & Xu, 2024) evaluate PLMMiiitializze decision transformers mainly on short-term environments such as MuloCo Locomotion and Atari. In contrast, we aim to investigate whether Markov heads also contribute next-action prediction in long-term environ- ments. To this end, we adopt LaMo (Shi et al., 2024), representative PLM-based decision transformer, the im- plementation of our baseline GPT2-DT, and refer to it as such throughout the rest of the paper. We then compare our proposed method, GPT2-DTMA, with several baselines, including the standard Decision Transformer (DT) (Chen et al., 2021), the value-based offline RL algorithm CQL
6

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
Table Mean normalized score (f) with three seeds MuJoCo Locomotion (short-term environment). The dataset names are abbreviated as follows: medium as medium- replay as m-r, medium expert m-e and Halfcht as Halfcheetah.
Dataset
CQL
DT
DC
DTMA
GPT2-DT
GPT2-DTMA
GPT2-P
Hopper-m
51.7
67. 17.4
79.7
65 65.242.3
77.9
77.4
70.8+14
Hopper-m-r
73.7
74.1+4.1
82.
74.7+3.7
77.9+2.1
80.4-1.8
75.3+2.9
Hopper-m-e
107-310.4
108.620.3
111.410.2
109.410.1
111.710.1
111.610.1
109.4
Walker2d-m
72.6+1.3
74.3+1.9
79.3+1.3
74.0+1.4
77.1
79.9=0.8
73.3+2.7
Walker2d-m-r
78.3
71.9422.
79.2+1.8
70.5+3..
74.0=2.9
77.0=2.5
71.543.5
Walker2d-m-e
107.8=0.2
107.6=0.1
108.3+0.2
108.1+0.2
108.1
108.640.1
108. 108.240.1
Halfcheetah-m
10.3
42.8
43.1
43.5.0.3
42.610.5
43.0
42.3
Halfrheetab-m-rr
42.4
39.2+2.4
38.343.4
38.843.3
39.8+3.1
40.3_2.9
39.6 39.643.1
Halfcheetah-m-ce
88.4+1.1
91.9+1.3
92.1+1.3
91.841.1
92.3+0.7
92.5-1.1
92.2+0.5
Table 4. Episode length (1) in different size of Maze (long-term environment).
Environment
Dataset
CQL
DT
DC
DTMA
GPT2-DT
GPT2-DTMA
umaze
58.0+1.0
56.3+1.7
54.3
58.7+1.3
59.043.0
55.042.
PointMaze
medium
134.0
158.717.3
236.3 16.7
155.0
188.0
146.3 146.319.3
large
162.0+6.0
195.3+14.7
288.0412.0
217.0+16.0
257.3+13.3
203.0+11.0
(Kumar al., 2020), and Decision Convformer (DC) (Kim al., 2023). Full experimental setup details are provided in Appendix C.
5.2. Main Results
To validate the effectiveness of MoA in regulating the at- tention heads of G G2T2-DT, evaluate the performance of GPT2-DTMA proposed in Section 4.2 across both short- term and long-term environments. As shown in Table 3, GP T2--TMM outperforms DT in most short-term environ- ments. Notably. MoA introduces only modest increase in model complexity approximately a 3% rise in param- eter count over the base DT-making a lightweight yet effecttve enhancement.
5.1. Short-term and Long-term Environments
In this section, provide brief introduction to each perimental task and categorize them into short-term and long-term environments, as follows:
MuJoCo Locomotion tasks (Fu al., 2020) are physics engines designed primarily for simulating and control- ling articulated rigid body systems. Kim et al. (2023) believe that short-term planning ability is enough for achieving a good performance in those tested environ- ments. a policy can control the agent to reach further or move faster while maintaining stable posture, then the normalized score will be higher.
In long-term environments, Table 4 shows that integrating the MoA architecture into GP GPT2 DT enables the agent to reach the goal in PointMaze with shorter episode lengths. However, as stated in Theorem 4.5, MoA can only regulate the influence of Markov heads, whose inherent properties are preserved through fine-tuning. Consequently, due to the adverse impact of Markov heads in long-term scenar- ios, GPT2-DTMA can reduce the performance gap with DT while suffering difficulty to outperform it. Detailed trends in how MoA adjusts the weights of Markov and Non-Markov heads across short-term and long-term environments are provided in Appendix
In a Maze task, such as PointMaze, policy model is trained to direct the object from a random start position to the goal in maze. Note that this point mass can only observe its own position and the position the goal; it has no knowledge of the positions of obstacles in the maze. The policy model needs learn how to outline the entire maze based on past trajectory information. Therefore, this environment is generally considered to require model's long-term planning ability (Lin et al., 2022). We evaluate whether the policy model can guide the point mass to the goal using a shorter episode length as the evaluation metrie.
By comparing the performance of the same model in Table 3 and Table 4, we observe that both DC and CQL exhibit strong performance only in specific scenarios that align with their respective planning ability biases, but struggle in other settings. In contrast, our model GPT2-DTMA with adaptive planning ability maintains relatively consistent and competitive performance across both short-term and long-

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
Table GPT2-DIMA with different context length (k) in PointMaze-large large. GMakow represents the sum of weights for all Markov heads. Ratatoo represents the percentage of GMackow rela- tive to that when k = 10.
Table 6. GP12-DIMA performance using different context length (k) in Hopper-medium environment. GMtatow represents the sum of weights for all Markov heads. Hatatoo represents the percentage
GMiikov relative to that when l = 10.
k
Episode length (1)
G Markov
Raarkov
k
Normalized score (1)
G Markov
Markow
:: 10
236.3
0.397
100.0%
k 10
77.1
0.505
100.0%
k 20
230.7
0.310
78.1%
= 20
77.4
0.488
96.6%
A 30
226.0
0.234
58.9%
k 30
76.3
0.480
95.0%
k 40
214.3
0.188
47.3%
k 40
77.0
0.461
91.3%
k= = 50
203.0
0.167
42.1%
k 50
75.8
0.440
87.1%
MoA allows Markov heads to retain a higher influence in action prediction, whereas in long-term environments, the model relies more heavily on other attention heads. These findings indicate that MoA can effectively identify whether the environment requires short-term or long-term planning ability and adaptively modulate weights across different heads suit the environment.
term environments.
Furthermore, we conduct an ablation study by incorporating the MoA mechanism into randomly initialized DT, referred to as DTMA. Table 3 shows that this method does not lead to performance gains, as evidenced by the comparable results between DTMA and DT. This suggests that the effectiveness of MoA primarily lies in its ability modulate the influence of Markov heads when predicting the next action, a property not present in models without pretraining.
6. Discussion
To further validate the critical role of Markov head in short- term environments, we design two different approaches: The first approach explicitly reduces the influence of Markov heads by lowering their MoA weights. The second approach modifies the initialization of the GPT2-DT model, using parameters that do not satisfy the Markov matrix property.
5.3. Attention Score Distribution of GPT2-DTMA in Different Environments
In this section, investigate whether GP GPT2-DTME is ca- pable of adaptively adjusting the weights of Markov heads to better adapt to different environments. We use the most challenging task, DointMaze-largee as our testbed, which a typical long-term environment. long-term environ- ments, models are required to extract useful information from past sequences, rather than relying primarily on current state. However, the Markov heads tend to focus attention disproportionately on the most recent input, limiting the model's 's ability to leverage Long-range context and thereby negatively affecting performance. By using MoA to reduce the influence of Markoy heads, GPT2 -DMA can alleviate their negative impact on action predictions. As shown in Table 5, it can be seen that the length of context length in- creases, both the sum of weight GMurkuv assigned Markov heads and the number of steps required for the point mass to reach the goal gradually decrease. This suggests that GPT2-DTMA endeavors to capture previous information that is further away from the current state.
6.1. Reducing the Weights of Markov Heads in Action Prediction
To assess the role of Markov heads in short-term envi- ronments, we conduct comparative experiment against G2T2-DM by intentionally reducing their contribution to action prediction. To this end, we develop a variant model called GPT2-P, which introduces a penalty on the weights of Markov heads in the loss function:
LOPT p = Top + a
Σ G((z),
(E(L5,100
where [1,5, 10] are the indices of Markov heads.
By applying a penalty coefficient of a = 0.1, the cumulative weight assigned to the Markov heads is reduced from 0.488 in GP D--T2-DMM
Similarly, we conducted a comparative experiment examin- ing the effect of context length in short-term environments. As context length increases, the model shows a growing ten- dency to attend to distant information, leading to an overall decrease in Garroo as shown in Table 6. However, at the same context length, the Rarkoy in short-term environments is significantly higher compared to those in long-tern envi- ronments. This suggests that, in short-term environments,
0.184 in its variant, GPT2-P. As shown in Table GPT2-P P fails match the policy performance of GPT2 -DT2A in short-term tasks. This degradation under- scores the significance of Markov heads in improving action selection under limited planning ability. In many cases, reducing their contribution causes GPT2-P to regress to the performance of the standard DT baseline. Furthermore, these results suggest that other components pretrained for
8

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
Table Mean normalized score (t) of DT initialized with different pretrained parameters in MuJoCo Locomotion.
Dataset
DT
GPT2-DT
CLIP-DT
GPTJ-DT
Hopper-m
67.4_1.7
77.9
66.911.3
71.611.5
Hopper-m-r
74.1-4.1
77.942.1
76.1
72.9+1.7
Hopper-m-e
108.6
111.710.1
109.210.2
109.640.1
Walker2d-m
74.3=1.9
77.1+1.1
74.2+0.8
74.8+1.2
Walker2d-m-r
71.9=2.7
74.0+22.9
70.3
70.4+2.3
Walker2d-m-e
107-610.1
108.140.1
107.940.3
108.1
Halfchectah-m
42.8=044
42.640.5
42.6=0.6
42.940.3
Halfcheetah-m-r
39.2=2.4
39.8
39.843.3
39.342.5
Halfcheetah-me
91.9=1.3
92.3+0.7
92.1
92.1+0.9
NLP the feed-forward network (FFN) lay- ers ma offer limited transferability to the reinforcement learning domain.
7. Conclusion
In this paper, addressed key question: how do pre- trained language models (PLMs) influence fine-tuning per- formance in offline reinforcement learning? We introduced the concept of Markov heads and investigated their distinc- tive properties. Our analysis reveals that Markov heads are the primary transferable inductive bias from PL.Ss, benefi- cial short-term tasks but detrimental long-term scenar- ios. To address this challenge, we proposed GPT2-D general and adaptive approach that enables the model to automatically adjust its planning capacity across differ- ent environments. Experimental results demonstrate that GP D2-MMA significantly narrows the performance gap in both short- and long-term settings. We believe that un- derstanding the mechanisms behind PLM transfer opens promising directions for future research, such optimiz- ing the of pretrained parameters in RL designing pretraining objectives tailored for decision-making tasks.
6.2. Initializing DT Using Pretrained Parameters without Markov Heads
In our offline RI. setup with an autoregressive formulation, we initialize the model using pretrained text encoders from CLIP and GPT-J. While CLIP includes both text and image encoders, we adopt only its text encoder, which is autore- gressive Transformer pretrained large-scale textual data. CLIP is trained to align textual descriptions with correspond- ing images, whereas GPT-J is GPT2-style causal language model, but pretrained on a different dataset (Pile) rather than the WebText corpus used by GPT2. Despite architec- tural similarity and pretraining on textual data, we find that the attention layers in both CLIP and GPT-J do not exhibit the Markov matrix structure observed in our CPT2-DTMA model. We attribute this to differences in training objectives: unlike GPT2, whose next-token prediction loss may encour- age local temporal dependencies, the contrastive objective in CLIP and the different pretraining distribution of GPT-J do not appear induce the same structural bias the attention heads.
A Acknowledgements
This work is partially supported by NSFC (61732006) and Public Computing Cloud, Renmin University of China.
As shown in Table 7, initializing DT with pretrained param- eters from CLIP (CLIP-DT) GPT-J (GPTJ-DT) does not yield performance comparable to GPT2-DT. The pres- ence of Markov heads in GPT2-D7 enables the model to allocate sufficient attention to the current state, resulting improved performance over the vanilla DT. The absence of the Markov mairix structure in CLIP-DT is expected, given its contrastive training objective. However, although GPT-J is pretrained using the same next-token prediction loss as GPT2, it fails to exhibit this property. This discrep- ancy may be attributed to differences in the pretraining data distribution and scale.
Impact Statement
This paper presents work whose goal is to advance the field of pretrained decision transformer. There are many potential socictal consequences of our work, none which we feel must be specifically highlighted here.
References
Badrinath, A., Flet-Berliac, Y., Nie, A., and Brunskill, Waypoint transformer: Reinforcement learning via su- pervised learning with intermediate targets. Advances in Neural Information Processing Systems, pp. 1-22, 2024.
Bau, D., Zhu, J.-Y., Strobelt, H., Lapedriza, A., Zhou, B.,
9

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
and Torralba, A. Understanding the role of individual units in deep neural network. Proceedings of the Na- tional Academy of Sciences, 117(48))30001-330777 2020.
Kostrikov, I., Nair, A., and Levine, S. Offline reinforce- ment learning with implicit q-learning. In International Conference Learning Representations, pp. 1-11, 2021.
Carroll, M., Paradise, O., Lin, J., Georgescu, R., Sun, M.,
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Con- servative q-learning for offline reinforcement learning. Advances in neural information processing systems, 33: 1 1179-1191, 2020.
Bignell, D., Milani, S., Hofmann, K., Hausknecht, M., Dragan, A., et al. Uni [mask]: Unified inference in se- quential decision problems. Advances in Neural Informa- tion Processing Systems, 35:35365-35378 2022.
Levine, S., Kumar, A., Tucker, G., and Fu, J. Offline rein- forcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0166 01645,,,,,,,,,, 2020.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34: 50841155-5 2021.
Lin, Q. Liu, H., and Sengupta, B. Switch trajec- tory transformer with distributional value approximation for multi-task reinforcement learning. arXiv preprint 0rXiv:2203.07413. 2022.
Correia, A. and Alexandre, L. A. Hierarchical decision transformer. In 2023 IEEF/RS/ International Conference on Intelligent Robots and Systems (IROS), pp. 1661-1666, 2023.
Littman, M. L. Memoryless policies: Thcoretical limita- tions and practical results. In From Animals Animats 3: Proceedings of the third international conference on simulation of adaptive behavior, volume 3. pp. 238. MIT Press Cambridge, MA, USA, 1994.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for lan- guage understanding. arXiv preprint arXiv:1810.04800 2018.
Liu, F., Liu, H., Grover, A., and Abbeel, P. Masked autoen- coding for scalable and generalizable decision making. Advances in Neural Information Processing Systems, 35: 2022.
Fu. J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072199 2020.
Mezghani, L., Bojanowski, P., Alahari, K., and Sukhbaatar, S. Think before you act: Unified policy for interleav- ing language reasoning with actions. arXiv preprint arXiv:2304. / pp. 1-9, 2023.
Furuta, H., Matsuo, Y., and Gu, S. S. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:21 10366, 2021.
Goel, R., Sulaiman, M., Noorbakhsh, K., Sharifi, M., Sharma, R., Jamshidi, P., and Roy, Pre-trained lan- guage transformers are universal image classifiers. arXiv preprint ar 1xXxx:2201.10188 2022.
Michel, P., Levy, O., and Neubig, G. Are sixteen heads really better than one? Advances neural information processing systems, 32, 2019.
Nair, A., Dalal, M., Gupta, A., and Levine, Accelerating online reinforcement learning with offline datasets. arXiv preprint wwXiv:2006.09352 pp. 1-17, 2020.
Janner, M., Li, Q., and Levine, S. Offline reinforcement learning one big sequence modeling problem. Ad- vances in neural information processing systems, 34: 1273-1286, 2021.
Noorbakhsh. K., Sulaiman, M., Sharifi, M., Roy, K., and Jamshidi, P. Pretrained language models symbolic mathematics solvers too! arXiv preprint arXiv:2110.03500 2021.
Jin, P., Zhu, B., Yuan, L., and Yan, S. Moh: Multi-head attention mixture-of-head attention. arXiv preprini ax1v:2410.11842. 2024.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, J. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of Machine Learning Research, pp. 1-67, 2020.
Kenton, J. D. M.-W. C. and Toutanova, L. K. Bert: Pre- training of deep bidirectional transformers for language understanding. In Proceedings of naacI--HI.T volume Minneapolis, Minnesota, 2019.
Kim, J., Lee, S., Kim, W., and Sung, Y. Decision con- vformer: Local filtering in metaformer is sufficient for decision making. arXiv preprint arXiv:2310.0000 2023.
Reid, M., Yamada, Y., and Gu, S. S. Can wikipedia help offline reinforcement learning? 1221 pp. 1-16, 2022.
arXiv preprint
10

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
Shang, J., Kahatapitiya, K., Li, X., and Ryoo, M. Star- former: Transformer with state-action-reward represen- tations for visual reinforcement learning. In European conference on computer vision, pp. 462-479, Springer, 2022.
Yamagata, T., Khalil, A., and Santos-Ro R. Q- learning decision transformer: Leveraging dynamic pro- gramming for conditional sequence modelling in offline rl. In International Conference on Machine Learning. pP. 38989-39007. PMLR, 2023.
Shi, R., Liu, Y., Ze, Y., Du, S. S., and Xu, H. Unleash- ing the power of pre-trained language models for offline reinforcement learning. In International Conference Learning Representations, pp. 1-24, 2024.
Yang. Y. and Xu, P. Pre-trained language models improve the few-shot prompt ability of decision transformer. arXiv preprint wrXiv:2408.014000 2024.
Yuan, L., Chen, Y., Wang, T., Yu, W., Shi, Y., Jiang, Z.- H., Tay, F. E., Feng, J., and Yan, S. Tokens-to-token vit: Training vision transformers from scratch imagenet. In Proceedings of the IEEE/CVF international conference on computer vision, pp. 558-567,200 2021.
Sun, Y., Ma, S., Madaan, R., Bonatti, R., Huang, F., and Kapoor, A. Smart: Self-supervised multi-task pretraining with control transformers. In International Conference on Learning Representations, pp. 1-21, 2022.
Zhang, X., Shen, Y., Huang, Z., Zhou, J., Rong, W., and Xiong, Z. Mixture of attention heads: Selecting attention heads per token. arXiv preprint arXiv:22 /0.0552/00051444 2022.
Takagi, On the effect of pre-training for transformer in different modality offline reinforcement learning. Advances in Neural Information Processing Systems, 35: 30678-30692, 2022.
Zheng, H., Shen, L., Luo, Y.. Liu, T., Shen, and Tao, D. Decomposed prompt decision transformer for efficient un- seen task generalization. Advances in Neural Information Processing Systems, 37: Net:122294-300000 2024.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.,, Gomez, A. N., Kaiser, L. u., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Zheng, Q., Zhang, A., and Grover, A. Online decision transformer. In International Conference on Machine Learning, pp. 27042-27055 2022.
Wang, T., Yang, Q. Wang. R., Sun, D., Li, J., Chen, Y., Hu, Y., Yang, C., Kimura, T., Kara, D., and Abdelza- her, T. F. Fine Fine-grained control of generative data aug- mentation in iot sensing. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https::////. http://ppeeeeeeeeeeeeeeeeee n forum? id=ZCygNDMITT.
Wu, P., Majumdar, A., Stone, K., Lin, Y., Mordatch, I,, Abbeel, P., and Rajeswaran, A. Masked trajectory mod- els for prediction, representation, and control. In Inter- national Conference on Machine Learning, pp. 37607- 37623, 2023.
Xie, 7., Lin, 7., Ye, D., Fu, Q., Wei, Y., and Li, S. Future- conditioned unsupervised pretraining for decision trans- former. In International Conference on Machine Learn- ing. pp. 38188-38823, 2023.
Xiong, S., Payani, A., Kompella, R., and Fekri, F. Large language models can learn temporal reasoning. In Pro- ceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 104700 2024.
Xu, M., Shen, Y., Zhang, S., Lu, Y., Zhao, D., Tenenbaum, J., and Gan, C. Prompting decision transformer for few- shot policy generalization. In International Conference on Machine Learning, pp. -2645,, 2022.

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
A. Pretrained Language Models for Offline RL Tasks
Autoregressive Language Model Pre-training
Fine-tuning on offline RL
Actions speak louder than words <eos>
at
Transformer
Decision Transformer
Token embedding layer
Return-to tooron State /Action embeedding layer
<bos>
Actions
speak louder than words
R
St
Figure 6. The paradigm transferring PLM parameters offline tasks,
Figure follows the paradigm for transferring knowledge from the natural language domain to the RI. domain in previous work (Reid al., 2022). In the first stage, we pretrain Transformer with language data to predict the next token. In the second stage, fine-tune pretrained Transformer downstream offline RL tasks to predict the next actions.
B. Comparisons of Attention Score Distributions
DT
GPT2-DT
Pewer Iterations
Walker266mediie
Hallcheotal menooaalmeediun mediium
Figure compare attention score distribution of and GPT2-DZ during the process of fine tine-tuning within fewer iterations more iterations in Walker2d-medium (upper row) and Halfcheetah-mmeium (lower row) environments.
We explore how the attention score distribution differs between DT and GPT2 DT of iterations versus a larger number of iterations. Figure 7 illustrates that in different environments, GPT2-DT can still leamn the ability to focus on closer positional information in fewer training iterations compared to DT. This ability to concentrate on recent information is precisely what needed to excel short-term environments.
DT when fine-tuning with smaller number
C. Experiment Setup
We GPT2-small pretrained parameters to initialize DT transformer layer. To fine-tune pretrained language models, set the learning rate le-4 with 10K warmup steps and weight decay le-4, We fine-tune the model with 100K steps and batch size as 64. Our training process uses single Nividia A100 and four Tesla V100 with cuda version of 12.3. During training, the agent is evaluated every 500 steps, by running 20 episodes. We report results averaged among 3 seeds (seeds 0, 0, 2 are used). The offline datasets used for training the model are collected in D4RL (Fu al., 2020) benchmark. For MuloCo Locomotion tasks, the tasks we selected are halfcheetah, hopper and walker2d. Each task has datasets collected
12

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
with different strategies: medium dataset consisting of 1M interaction samples of a middle-level RL agent. medium-replay includes the whole replay buffer when training an RI. agent until it reaches middle-level performance. medium-expert is mixture of the medium dataset and IM interaction samples of expert-level RL policy. For Maze tasks, umaze, medium and large represent square mazes with side lengths of 5, 8, and 12, respectively.
D. Influence of Markov heads in Long-term Environments
3.4
912
Figure 8. Importance different heads in PointMaze (long-term environment). The red bar represents the importance scores allocated to the Markov head in the first attention layer.
Figure 8 shows the importance scores of different attention heads when predicting actions. Similar to the phenomenon in short-term environments, the Markov heads exhibit large influences predicting actions. It indicates that the poor performance of GPT2-DDD DD long-term environments may stem from the Markov heads focusing only on the current state.
E. Theoretical Proofs
In this section, we present proofs of Theorem 4.3, Corollary 4.4 and Theorem 4.5.
E.1. Proof of Theorem 4.3
We will give simple proof showing that Markov matrix, i.e., W?W satisfies Definition 4.1, will pay most attention
the last-input token under random embedding initializations, i.c., E E(W"W;"" E also satisfies Definition 4.1, where
E E R Kxd is random embedding matrix, K is the sequence length and in this case, A = E E(W?W E
Proof. For any em = (et,, 2n ... ed and cn (e en e
ed). where m, n (1, K) and + n, we have
E El 1
Aij Aly ell
i= 1=
A E. m e +
Σ
Alenee
(10)
d
Ale + + (Elein))
Σ
Aid
d
Σ
Ele ot 2=]=0 = 00
1=1
where the last equality is because E(ee.]] B B[e
0.
And we know that ef, and et, are i.i.d, hence by taking the
13

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
expectation, we obtain that
(Elctn)) 2 < E (2
E(e - [[en
(11)
where the last inequality is due to Varleml Equation (10), we have
B (e
(Blend)
> 0. Also, by simply replacing e with e in
E
pm
Ae
Σ
A (e/m)) + (Fleml) 2
Σ
Asj
i=1
i,j=1," dditj
(12)
Σ
A
(e/n))
Σ A
Since we know that {4,,)1=1 are all positive, we can conclude that (e > B(em 0. Therefore, for any > Definition 4.1 will be satisfied.
E.2. Proof of Corollary 4.4
Proof. By Theorem 4.3, we know that E[O] We apply concentration bounds for sub-exponential random variables, noting that [I, is sub-exponenties with sub- Gaussian proxy || ||]| For the diagonal mean D. there are K i.i.d. terms, and for the off-diagonal mean O, there are K K(K' - 1) terms. By Bernstein inequality and Wriigt inequality, we obtain:
E|D
> r for some > 0. And for any E > 0, we have E|D : > r.
2r
P(D R[D] 2 E) S 2exp (-cs K min
E2
))
(13)
||||
and
E² E
)
P (IO [[O]] 2 +) 2xx -ogK 1) min
(14)
[TIIT |
Taking union bound and noting that if 2 D & and
< E + then
0 > E0 +
E -
(15)
Then we have
D
PP <r) <P(D < E -) + P(O E[O] + 8o.
(16)
where do A 2 exp (-e1 K min of E, we obtain the desired high-probabiln bound.
2 | T
+ exp
(-o2K(K 1) min
Tai T
c2 Therefore under the definition
E.3. Proof of Theorem 4.5
In this section, we will discuss the phenomenon in Table. 2 and give theoretical analysis for Theorem 4.5.
Proof. Suppose that the total number of training iterations is K, and A w'W** is Markov matrix, the difference between the initialized parameters A and the trained parameters denoted as AAK = Σk Tegk, where Tik denotes the learning rate the k-tth iteration and A A AA", By choosing the Adam optimizer, know that there exists No such that 7)& & Nio for all l =1... 1,.... K. Hence, have
K
Oaxd 5 AAA =
Σ
nhg <
Th gkl
S m7 KoK B laxd-
(17)
km
14

---

Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning
Let Omaz
max [ sylwwhree j -,,... ,dd and amaz S %ooKK (BB. B. Due to Definition 4.1, we know that there exists ,...,,, 2 r p v, le., AW 2 (r + p) p) 4/. Since amar m .d Ale. then we have
p such that 1
A = A AA 2 A Omar > 0,
(18)
|A2y1 amar 1 A [4
AAA
S [Ag]]
+ amar-
thus, 1 2 A3j Omaz and max max{1A]) [A
Omaz: 0) -
S 1 + amar. Since Ounx S NoKB < r41 then
(r + 1)amaz < $l < 4 r- 14.yl. i.c,,
A
4
amav
2
r.
(19)
14
+ amax
Hence, AK is still Markov matrix.
F. FFN Cannot Achieve Adaptive Attention Facing Different Environments
Dataset GPT2-DT
GPT2-DT (part)
GPT2-DTMA
umaze
59
67
55
medium
188
229
146
large
257
286
203
Table Episode length (4) in different size PointMaze (a long-term environment). GPT2-DT fine-tune all parameters the model and GPT2-DT (part) only fine-tune the parameters in embedding layer and FFN layer.
The FFN layer can weight the output information from different attention heads, but unlike MoA, which assigns a shared importance score to the output of each attention head, the FFN applies nonlinear transformation to the outputs of the attention heads. To verify the necessity of our proposed combining GPT2-DT with MoA, it can be observed from Table that relying solely on the FFN layer is insufficient for effectively balancing the influences of the Markov heads and Non-Markov heads in action prediction.
G. Changes of MoA Weight between Markov Heads and Non-Markov Heads
wight
014
322-
Mrn on,, Alea mym
3.0
Checog
Figure 9. We illustrate the changes of MoA weight hetween Markov heads and Non-Markov heads in short-term environment (left) and long-term environment (right).
To better observe the differences in the impact of MoA in short-term and long-term environments, we show the MoA weights assigned for Markov heads and Non-Markov heads while training in Figure 9. We found that, after the MoA weights have converged, the Markov heads is more influential than the Non-Markov heads in predicting actions in short-term environments, while in long-term environments, the Non-Markov heads hold greater importance. This indicates that combining the MoA structure enables the model identify whether the planning ability required by the environment is short-term or long-term.
15


Reviews:
[{"review_id": "XJ6h4wKYWM", "reviewer": "Reviewer_WJgz", "summary": "This paper investigates why pre-trained language models (PLMs) boost Decision Transformer performance in offline RL setting. The authors identify crucial \"Markov heads\" within PLMs that strongly focus attention on the most recent input state. While beneficial for short-term tasks like MuJoCo, theoretical analysis and experiments show these heads are rigid and cannot be easily changed via fine-tuning, thus hindering performance in long-term planning tasks like Mazes. \n\nTo address this limitation, the paper introduces GPT-DTMA, which uses a Mixture of Attention mechanism to adaptively weight the influence of different attention heads. This allows the model to dynamically control the impact of Markov heads based on the specific environment's requirements. Results demonstrate that GPT-DTMA outperforms baselines in short-term settings and significantly closes the performance gap in long-term scenarios.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable comments. Please kindly find the response to your concerns below.\n\n**W1. The improvement of GPT-DTMA is within the confidence interval.**  \nThe 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.\n\n**W2. The paper doesn't extensively compare its absolute performance against the broader state-of-the-art in offline RL methods on these benchmarks.**  \nWe 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. \n\n**Table R2**:\n| Dataset (short-term) | CQL | DT | GPT-DT | DC | GPT-DTMA |\n| --- | --- | --- | --- | --- | --- |\n| Hopper-m | 58.1 | 67.4 | 77.9 | **79.5** | 77 |\n| Hopper-m-r | 75.3 | 74.1 | 77.9 | **82.1** | 80.4 |\n| Walker2d-m | 72.7 | 74.3 | 77.1 | 79.3 | **79.9** |\n| Walker2d-m-r | 78.6 | 71.9 | 74.0 | **79.1** | 77.0 |\n\n**Table R3**:\n| Dataset (long-term) | CQL | DT | GPT-DT | DC | GPT-DTMA |\n| --- | --- | --- | --- | --- | --- |\n| PointMaze-large | **167.7** | 195.3 | 257.3 | 276 | 203.0 |\n\n**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?**  \nWe 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. \n\n[1] Chen, Mark, et al. \"Generative pretraining from pixels.\" International conference on machine learning. PMLR, 2020.\n\n**Table R1**:\n| Dataset (short-term) | GPT-DT | GPTJ-DT | ImageGPT-DT |\n| --- | --- | --- | --- |\n| Hopper-m | 77.9 | 72.5 | 7.3 |\n| Hopper-m-r | 77.9 | 73.8 | 7.6 |\n| Walker2d-m | 77.1 | 75 | 1.2 |\n| Walker2d-m-r | 74.0 | 70.3 | 12.3 |\n\n**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?**  \nWe 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)$.\n\nThank again for your thoughtful reviews. We will incorporate the discussions in the revised version and we are looking forward to further discussions.", "claims_and_evidence": "The paper does a good job supporting the claims. \n\nFor experimental evidence, the main issue is the improvement of GPT-DTMA—for most cases in Table 3, the improvement is within the confidence interval, which weakens the validity of the proposed algorithm. Even for long-term scenario in table 4, the case is not that better.", "methods_and_evaluation_criteria": "The proposed methods overall make sense. For the benchmark environment, do you have any other experiments beyond the robot environment or maze environment? They are relatively old environments in the offline RL setting.", "theoretical_claims": "I didn't check the details but the overall logics seem fine. My main point is that the author should address more on how some of the assumptions/conditions are determined. For example, for more explanation on the physical meaning of definition 4.1 and 4.2 can make it more reader-friendly.", "experimental_designs_or_analyses": "I have checked the validity of experiments, and most are reasonable, clear. My only issue is the significance of table 3/4.", "relation_to_broader_scientific_literature": "It is closely related with the LLM and pretraining literature. The author properly discuss its relationship in the related work and preliminary section.", "essential_references_not_discussed": "All references are discussed.", "other_strengths_and_weaknesses": "Strengths:\nThe overall idea is interesting and the demonstration process is quite clear and sound. The proposed Markov Head Concept is novel and properly addressed.\n\nweakness:\nWhile GPT-DTMA shows improvements over GPT-DT and DT baselines, the paper doesn't extensively compare its absolute performance against the broader state-of-the-art in offline RL methods on these benchmarks (which might include non-Transformer methods or different Transformer variants). The focus is more internal to understanding PLM transfer within the DT framework.", "other_comments_or_suggestions": "No", "questions_for_authors": "1. 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?\n2.  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?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ypu3ygFQtO", "reviewer": "Reviewer_vRU6", "summary": "Incorporating Pretrained Language Models (PLMs) into Decision Transformers (DTs) has shown promise in the area of offline reinforcement learning (RL). However, it is unclear why the representations obtained from NLP tasks would be beneficial for RL tasks. The authors aim to address this question by analyzing the attention heads of several models and demonstrating that they exhibit Markov properties. Identifying whether these heads are Markov or not is crucial for understanding the limitations of these models in solving short-term and long-term tasks in various environments. \n\nThe authors perform several experiments in both short- and long-term environments and show the impact of Markov heads on model performance. They also conduct various ablation studies to further support their findings, such as examining the relationship between context length and Markov head weight, and initially down-weighting Markov heads to see how this directly affects performance.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your valuable comments. Please kindly find the response to your concerns below.\n\n**Q1. Would a MoE (mixture of experts) perform better in these long-term environments?**  \nIt’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. \n\n**Q2. What is the high-level intuition behind a Markov head? Additionally, why is diagonal dominance important in the context of NLP and RL?**  \nThe 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.\n\n**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?**  \nFollowing 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.\n\n**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?**  \nNon-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.\n\n**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.**  \nWe 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. \n\n**Q6. Why do you need to adaptively control the weight of the attention heads?**  \nDue 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.\n\n**Q7. Do you have a skyline model that shows what good behavior looks like?**  \nWhile 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.\n\n**Q8. How does the Markov heads change for DT and GPT-DT as in Table 5/6?**  \nWe 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.\n\n**Q9. With an increase in context length would the models tend to focus more easily on distant information?**  \nIn 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.\n\nThank you again for your reviewing efforts and we sincerely hope our responses have addressed your concerns.", "claims_and_evidence": "The authors claim that PLMs influence performance in offline RL based on Markov heads, which are the key information transferred from PLMs. They argue that this is beneficial only for short-term environments and has a negative impact on long-term environments. The evidence that the authors show is based on empirical observation that supports their theory. \n\nThe evidence is not entirely convincing; there has been a lot of research on transformer attention regarding the significance of each head and the observation that not all attention heads are necessary. However, this paper seems disconnected from that body of work. Additionally, it is unclear whether the issue the authors are addressing is related to a distribution shift or a modeling problem, as the paper does not clarify how the short- and long-term data splits were determined. It also does not specify whether the model was trained and tested in the same or different environments.", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria make sense for the problem at hand.", "theoretical_claims": "No, I did not check the correctness of the proofs in the paper.", "experimental_designs_or_analyses": "The author conducted several experiments and analyses to support the claims in their paper:\n\n- Main Results: The main experiments were performed on both short-term tasks (using Mujoco) and long-term tasks (using PointMaze and AntMaze). The authors compared their results to several baselines, which were trained using PLM and from scratch. They analyzed their findings and briefly explained that their proposed algorithm outperforms the others.\n\n- Ablation Studies: The primary focus of the authors’ ablation studies was to compare the relationship between context length and Markov head weights in both short-term and long-term tasks. Although they demonstrated that Markov head weights are higher for short-term than long-term tasks, the paper does not include sample ablations for other models, making it difficult to fully understand the differences.", "relation_to_broader_scientific_literature": "The key contribution to the broader scientific literature is understanding when the representations of PLMs will be helpful in solving RL tasks. In particular, the papers suggest that PLMs learn what are referred to as Markov Heads, which RL tasks can take advantage of. However, these Markov Heads may only be adequate for short-term tasks and not for long-term tasks. Therefore, the broader scientific connection is to determine when PLMs will be beneficial for solving RL tasks.", "essential_references_not_discussed": "- MOH: MULTI-HEAD ATTENTION AS MIXTURE-OF-HEAD ATTENTION by Jin et al. 2024\n- Are Sixteen Heads Really Better than One?  by Michel et al. 2019", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Would a MoE (mixture of experts) perform better in these long-term environments?\n\n2. What is the high-level intuition behind a Markov head? Additionally, why is diagonal dominance important in the context of NLP and RL? I'm a bit confused about the significance of having an attention head that exhibits diagonal dominance. I would assume some attention heads display diagonal dominance while others do not, as each attention head focuses on different patterns within the data. \n\n3. 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?\n\n4. Not all attention heads are Markov heads. 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?\n\n5. Section 4.2: \"In real applications, it is hard to determine whether the planning ability is required by the current environment without prior knowledge\". I am uncertain whether the issues discussed stem from distribution shift problems due to changes in the training and testing environments or if they are related to modeling flaws.\n\n6. Why do you need to adaptively control the weight of the attention heads? It does not seem like you are addressing the modeling problem.\n\n7. Do you have a skyline model that shows what good behavior looks like, even if it's achieved through cheating? For example, imagine training a model on certain tasks and then testing it on different tasks. A skyline model would be trained on the test tasks to observe the behavior of the model as if it had seen the training datasets.\n\n8. How does the Markov heads change for DT and GPT-DT similar to the experiments shown in table 5 and table 6?\n\n9. I assume that with an increase in context length, regardless of the model, the models will tend to focus more easily on distant information because there is more content to consider. This seems like a fairly general phenomenon that is not specific to GPT-DTMA.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "53MikZXj8H", "reviewer": "Reviewer_cgss", "summary": "The paper identifies Markov heads in Pretrained Language Models (PLMs), attention heads with extreme focus on the most recent token. These heads transfer to Decision Transformers (DTs) in offline Reinforcement Learning (RL), improving short-term (Markovian) tasks but harming long-term planning. The paper introduces GPT-DTMA, which uses a Mixture of Attention (MoA) mechanism to adaptively weigh heads, balancing short-term and long-term performance. Theoretical analysis proves Markov heads persist under fine-tuning and experiments confirm GPT-DTMA improves performance across various RL tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**S1&Q3. Test other PLMs to confirm generality.**  \nWe 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.\n\n**Table R1**:  \n|Dataset(short-term)|GPT-DT|GPTJ-DT|ImageGPT-DT|\n|-|-|-|-|\n|Hopper-m|77.9|72.5|7.3|\n|Hopper-m-r|77.9|73.8|7.6|\n|Walker2d-m|77.1|75|1.2|\n|Walker2d-m-r|74.0|70.3|12.3|\n \n[1] Chen, Mark, et al. \"Generative pretraining from pixels.\" *International conference on machine learning*. PMLR, 2020.\n\n**S2. Analyze MoA behavior per time step (how weights change dynamically).**  \nWe analyzed the step-wise changes of average weights for Markov heads and non-Markov heads. Fig.11 and Fig.12 in [[URL](https://anonymous.4open.science/r/submission9905-7D09/readme.md)] 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.\n\n**S3. Compare GPT-DTMA’s performance against value-based offline RL baselines.**  \nWe 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. \n\n**Table R2**:\n|Dataset(short-term)|CQL|DT|GPT-DT|DC|GPT-DTMA|\n|-|-|-|-|-|-|\n|Hopper-m|58.1|67.4|77.9|**79.5**|77|\n|Hopper-m-r|75.3|74.1|77.9|**82.1**|80.4|\n|Walker2d-m|72.7|74.3|77.1|79.3|**79.9**|\n|Walker2d-m-r|78.6|71.9|74.0|**79.1**|77.0|\n\n**Table R3**:\n| Dataset (long-term) | CQL | DT | GPT-DT | DC | GPT-DTMA |\n| --- | --- | --- | --- | --- | --- |\n| PointMaze-large | **167.7** | 195.3 | 257.3 | 276 | 203.0 |\n\n**S4. Extra related works**  \nThanks 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.\n\n**Q1. How was the threshold for Markov heads chosen?**  \nWe 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.\n\n**Q2. Would a manual reduction of Markov head influence improve long-term tasks further?**  \nA 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. \n\n**Table R4**:\n| Dataset (long-term) | DT | GPT-DT | GPT-DTMA | GPT-DT-ALiBi-R |\n| --- | --- | --- | --- | --- |\n| PointMaze-large | 195.3 | 257.3 | 203.0 | 216 |\n\nFor the minor clarity issues, we will revise and address this issues accordingly. \n\nThank you again for your constructive comments!", "claims_and_evidence": "I find the claims supported by convincing evidence.\n\n1. **Markov heads exist in PLMs and aid short-term RL.**\n- The importance score shows GPT-DT's key heads have diagonal-dominant weight matrices.\n- GPT-DT outperforms standard DT in short-horizon MuJoCo tasks.\n- CLIP-DT lacks Markov heads and performs worse, confirming their significance.\n\n2. **Markov heads hurt long-term planning**\n- GPT-DT underperforms on PointMaze/AntMaze (long-horizon), taking more steps.\n- GPT-DTMA's MoA down-weights Markov heads, improving performance in long tasks.\n\n3. Markov heads persist through fine-tuning\n- Theorems show that Markov matrices remain dominant after small-gradient updates\n- Empirical results confirm that attention remains last-token-focused post finetuning.", "methods_and_evaluation_criteria": "I find the methods and evaluation criteria make sense.\n\n### Methods\n- Defines Markov heads as diagonal-dominant attention heads.\n- Identifies Markov heads by measuring zeroed-out head importance.\n- GPT-DTMA introduces gated MoA, letting the model adapt head weights dynamically.\n\n### Evaluation\n- Short-horizon tasks: Standard offline MuJoCo locomotion tasks (normalized return)\n- Long-horizon tasks: Standard offline Ant-Maze, Point-Maze tasks (steps to goal)\n- Baselines: DT, GPT-DT, DTMA, GPT-DTMA, CLIP-DT, GPT-DTMA-R (penalizing Markov heads)\n- Robustness: mean and standard deviation over multiple seeds", "theoretical_claims": "This paper is not a fully empirical paper. To support their main claim about Markov heads, this paper provides two theorems, which I found reasonable with correct proofs.\n\n1. Theorem 4.3: Markov heads retain their property under random projection (embedding changes).\n\n2. Theorem 4.5: Markov heads remain stable under bounded-gradient finetuning.", "experimental_designs_or_analyses": "The experimental design and analyses seem valid. The only thing I would suggest is to make comparisons to non DT-based offline RL methods (e.g., Decision Convformer or value-based methods)\n\n### Datasets\n- MuJoCo (short-term)\n- Point/Ant Maze (long-term)\n\n### Key findings\n- GPT-DTMA matches GPT-DT in short-horizon tasks while significantly improving long-horizon tasks.\n- GPT-DTMA-R degrades in short-term, confirming Markov heads aid short-horizon.\n- CLIP-DT performs worse than GPT-DT, proving Markov heads are beneficial.\n\n### Baselines\n- I like that the authors choice of baselines to control pertraining (GPT-DT), architecture (DTMA), and alternative PLMs (CLIP-DT)\n- I suggest the authors to make a comparison to Decision Convformer (Kim et al., 2024), since the paper also discusses Markov properties of short-horizon offline RL tasks.\n- It would be great if the authors could make comparisons to value-based offline RL methods (e.g., CQL, IQL) since they outperform Transformer based methods on sub-optimal datasets.", "relation_to_broader_scientific_literature": "This work is related to broader scientific literature and contributes as follows.\n\n- Background: Decision Transformers, PLM-based RL\n- Novel insight: identifies Markov heads as the mechanism behind PLM gains in RL.\n- Unique approach: MoA provides adaptive attention, unlike past methods relying on trajectory planning (Trajectory Transformer) or hierarchical policies.\n- Contribution to NLP: connects transformer interpretability (head specialization) with RL decision-making.", "essential_references_not_discussed": "This paper discusses most related references. One I may suggest is\n\n1. **Trajectory Transformer** \nJanner et al., Offline Reinforcement Learning as One Big Sequence Modeling Problem, 2021, for long-horizon transformer planning", "other_strengths_and_weaknesses": "### Strengths\n\n- Original: First work identifying Markov heads and their role in RL.\n\n- Significant: Improves transformer-based RL adaptivity across different timescales.\n\n- Clarity: Good structure; clear motivation and results.\n\n- Balanced theory & practice: Strong theoretical foundation + empirical validation.\n\n### Weaknesses\n\n- Limited PLM scope: Only tested on GPT-2 small; unclear if Markov heads exist in other PLMs.\n\n- Complexity: MoA adds parameters, though minimal overhead.\n\n- No direct comparison to RL SOTA: Does not benchmark against CQL/IQL in D4RL.\n\n- Minor clarity issues: Typos and vague phrasing (e.g., \"unknown-term environments\").", "other_comments_or_suggestions": "I enjoyed this work, and some minor suggestions that can improve this work, or some follow-up work ideas I can think of are as follows.\n\n1. Test other PLMs to confirm generality.\n\n2. Analyze MoA behavior per time step (how weights change dynamically).\n\n3. Compare GPT-DTMA’s performance against value-based offline RL baselines (CQL/IQL).", "questions_for_authors": "Here are some questions that I found confusing.\n\n1. How was the threshold for Markov heads chosen?\n\n2. Would a manual reduction of Markov head influence improve long-term tasks further?\n\n3. Have you tested Markov heads in other PLMs?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qvGECUzxDV", "reviewer": "Reviewer_L8Sr", "summary": "Previous works in the area of reinceforcement learning (RL)/foundation models have shown that pre-trained language models (PLMs) can enhance the performance of offline RL. This paper studies an important question: what kind of knowledge from PLMs has been transferred to RL to achieve such good results? They study the attention score distribution of decision-transformer (DT) and PLM-DT, and find that after training some iterations, the score distribution of DT is similar to the initial distribution of PLM-DT. They thus defined the concept of Markov heads, and further show that Markov heads are the key property transferred from PLM, both theoretically and empirically. Furthermore, they find that Markov heads are only beneficial for short-term environments, empirically. And they propose a general approach called GPT-DPMA, which show advantages over GPT-DT across both long-term/short-term tasks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**W1. For figure 1 and figure 4, DT is tested on only one task.**  \nWe have tested DT and GPT-DT on more tasks and the results are shown in Fig.1~10 at [[URL](https://anonymous.4open.science/r/submission9905-7D09/readme.md)]. The conclusion are consistent among different tasks. We will include these results from multiple tasks in the final version. \n\n**W2. Stronger results for Theorem 4.3**  \nWe 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}|, \nO \\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.\n\n**W3. Explanations for the scale**$K$  \nRecall 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. \n\n**Q1. In long-term task, why GPT-DTMA cannot surpass DT?**  \nIn 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.\n\n**Q2. Would the positional embedding play a role in yielding Markov heads?**  \nWe 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.\n\nThanks again for your constructive suggestions and we will include them in the final version.  We are looking forward to further discussions.", "claims_and_evidence": "**claim 1**: the score distribution of DT is similar to the initial distribution of PLM-DT\n\nevidence: Figure 1, DT is only tested on one task, not enough\n\n**claim 2**: the markov heads will not change after fine-tuning\n\nevidence: Figure 4, only tested on one task, not enough\n\n**claim 3**: markov heads are beneficial for short-term tasks, while not for long-term tasks\n\nevidence: Table 3, solid\n\n**claim 4**: MoA can reduce the performance gap in long-term tasks\n\nevidence: Table 4, solid", "methods_and_evaluation_criteria": "Yes, the proposed methods are practical, and the benchmark d4rl is commonly used.", "theoretical_claims": "I've checked the proofs for all theoretical results (Theorem 4.3 and Theorem 4.5). \n\nThough the proofs are correct, the results are not very strong. 1. The result that $\\mathbb E[EAE^\\top]$ is a markov matrix is not very useful. It would be better if the authors could prove that, w.h.p., $EAE^\\top$ is a markov matrix. 2. The scale of $K$ is uncertain, and might be very small.", "experimental_designs_or_analyses": "I've checked the experimental parts. The experiments are conducted in a common way in this area. One issue is that some of the experiments are only conducted on one specific task, like hopper. More systematic experiments should be conducted.", "relation_to_broader_scientific_literature": "I think this paper is important for this area. Many prior works [1,2,3] show the power of PLMs for downstream low-level RL tasks, however, this paper would be the first to offer a reasonable explanation for this phenomenon. Therefore, it would be of significant value for the study of foundation models for decision-making, if their claims are further systematically supported.\n\n[1] Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning. Shi et al. ICLR 2024.\n\n[2] Decomposed Prompt Decision Transformer for Efficient Unseen Task Generalization. Zheng et al. NeurIPS 2024.\n\n[3] Pre-trained Language Models Improve the Few-shot Prompt Ability of Decision Transformer. Yang et al. arXiv 2024.", "essential_references_not_discussed": "I don't think there is any work essentially related to the work not cited.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is well-written.\n- The observation is novel.\n\nWeakness:\n- See my comments regarding to ``Claims And Evidence`` and ``Theoretical claims``.", "other_comments_or_suggestions": "no", "questions_for_authors": "- On long-term tasks (Table 4), GPT-DTMA still cannot surpass DT. Could the author give some explanations?\n- Would the positional embedding play a role in yielding Markov heads? For example, if the positional embedding is relative, then the model might be trained to attend to nearest token.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper identifies that Markov heads in a PLM can transfer to a DT in offline RL, improving performances on short-term environments while degrading performances on long-term environments. Based on this observation, it proposes the mixture of attention (MoA) mechanism to adaptively weight different attention heads so that it especially down-weights Markov heads for long-term planning tasks. The empirical and theoretical analysis on the relationship between the existence of Markov heads and RL performances of PLM-based DT would be interesting, and the proposed MoA would be technically sound with sufficient empirical validation. Most concerns raised by the reviewers are addressed by the authors. Based on the consensus among the reviewers, I would recommend this paper to be accepted. 

I think more extensive empirical validation with recent small language models and on diverse real RL tasks would be helpful for this paper to be more valuable. For example, the performance gaps for short-term environments between GPT-DT and GPTJ-DT seems to be small even though there is no Markov head in GPTJ-DT.
Decision: Accept (poster)