SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Haoran He  $^{1,2*}$  Chenjia Bai  $^{2\dagger}$  Kang Xu  $^{2,3}$  Zhuoran Yang  $^{4}$  Weinan Zhang
Dong Wang² Bin Zhao²,5 Xuelong Li²,5†
$^{1}$ Shanghai Jiao Tong University  $^{2}$ Shanghai Artificial Intelligence Laboratory  $^{3}$ Fudan University  $^{4}$ Yale University  $^{5}$ Northwestern Polytechnical University
Diffusion models have demonstrated highly-expressive generative capabilities in vision and NLP. Recent studies in reinforcement learning (RL) have shown that diffusion models are also powerful in modeling complex policies or trajectories in offline datasets. However, these works have been limited to single-task settings where a generalist agent capable of addressing multi-task predicaments is absent. In this paper, we aim to investigate the effectiveness of a single diffusion model in modeling large-scale multi-task offline data, which can be challenging due to diverse and multimodal data distribution. Specifically, we propose Multi-Task Diffusion Model (MTDIFF), a diffusion-based method that incorporates Transformer backbones and prompt learning for generative planning and data synthesis in multi-task offline settings. MTDIFF leverages vast amounts of knowledge available in multi-task data and performs implicit knowledge sharing among tasks. For generative planning, we find MTDIFF outperforms state-of-the-art algorithms across 50 tasks on Meta-World and 8 maps on Maze2D. For data synthesis, MTDIFF generates high-quality data for testing tasks given a single demonstration as a prompt, which enhances the low-quality datasets for even unseen tasks.
The high-capacity generative models trained on large, diverse datasets have demonstrated remarkable success across vision and language tasks. An impressive and even preternatural ability of these models, e.g. large language models (LLMs), is that the learned model can generalize among different tasks by simply conditioning the model on instructions or prompts [49, 46, 7, 12, 60, 42, 66]. The success of LLMs and vision models inspires us to utilize the recent generative model to learn from large-scale offline datasets that include multiple tasks for generalized decision-making in reinforcement learning (RL). Thus far, recent attempts in offline decision-making take advantage of the generative capability of diffusion models [54, 20] to improve long-term planning [21, 2] or enhance the expressiveness of policies [63, 11, 8]. However, these works are limited to small-scale datasets and single-task settings where broad generalization and general-purpose policies are not expected. In multi-task offline RL which considers learning a single model to solve multi-task problems, the dataset often contains noisy, multimodal, and long-horizon trajectories collected by various policies across tasks and with various qualities, which makes it more challenging to learn policies with broad generalization and transferable capabilities. Gato [47] and other generalized agents [29, 65] take transformer-based architecture
*The work was conducted during the internship of Haoran He at Shanghai Artificial Intelligence Laboratory.
†Corresponding authors: Chenjia Bai (baichenjia@pjlab.org.cn), Xuelong Li
Figure 1: Overall architecture of MTDIFF. Different colors represent different tasks.  $S$ ,  $A$  and  $R$  denote the state sequence, action sequence, and reward sequence from multi-task data, respectively.  $S_{\mathrm{prev}}$  and  $R_{\tau}$  represent historical states and normalized return.
[62] via sequence modeling to solve multi-task problems, while they are highly dependent on the optimality of the datasets and are expensive to train due to the huge number of parameters.
To address the above challenges, we propose a novel diffusion model to further explore its generalizability in a multi-task setting. We formulate the learning process from multi-task data as a denoising problem, which benefits the modeling of multimodal data. Meanwhile, we develop a relatively lightweight architecture by using a GPT backbone [45] to model sequential trajectories, which has less computation burden and improved sequential modeling capability than previous U-Net-based [48] diffusion models [21, 33]. To disambiguate tasks during training and inference, instead of providing e.g. one-hot task identifiers, we leverage demonstrations as prompt conditioning, which exploits the few-shot abilities of agents [50, 64, 69].
We name our method the Multi-Task Diffusion Model (MTDIFF). As shown in Figure 1, we investigate two variants of MTDIFF for planning and data synthesis to further exploit the utility of diffusion models, denoted as MTDIFF-P and MTDIFF-S, respectively. (a) For planning, MTDIFF-P learns a prompt embedding to extract the task-relevant representation, and then concatenates the embedding with the trajectory's normalized return and historical states as the conditions of the model. During training, MTDIFF-P learns to predict the corresponding future action sequence given the conditions, and we call this process generative planning [78]. During inference, given few-shot prompts and the desired return, MTDIFF-P tries to denoise out the optimal action sequence starting from the current state. Surprisingly, MTDIFF-P can adapt to unseen tasks given well-constructed prompts that contain task information. (b) By slightly changing the inputs and training strategy, we can unlock the abilities of our diffusion model for data synthesis. Our insight is that the diffusion model, which compresses informative multi-task knowledge well, is more effective and generalist than previous methods that only utilize single-task data for augmentation [38, 52, 28]. Specifically, MTDIFF-S learns to estimate the joint conditional distribution of the full transitions that contain states, actions, and rewards based on the task-oriented prompt. Different from MTDIFF-P, MTDIFF-S learns to synthesize data from the underlying dynamic environments for each task. Thus MTDIFF-S only needs prompt conditioning to identify tasks. We empirically find that MTDIFF-S synthesizes high-fidelity data for multiple tasks, including both seen and unseen ones, which can be further utilized for data augmentation to expand the offline dataset and enhance policy performance.
To summarize, MTDIFF is a diffusion-based method that leverages the multimodal generative ability of diffusion models, the sequential modeling capability of GPT architecture, and the few-shot generalizability of prompt learning for multi-task RL. To the best of our knowledge, we are the first to achieve both effective planning and data synthesis for multi-task RL via diffusion models. Our contributions include: (i) we propose MTDIFF, a novel GPT-based diffusion model that illustrates the supreme effectiveness in multi-task trajectory modeling for both planning and data synthesis; (ii) we incorporate prompt learning into the diffusion framework to learn to generalize across different tasks and even adapt to unseen tasks; (iii) our experiments on Meta-World and Maze2D benchmarks demonstrate that MTDIFF is an effective planner to solve the multi-task problem, and also a powerful data synthesizer to augment offline datasets in the seen or unseen tasks.
2.1 Reinforcement Learning
MDP and Multi-task MDP. A Markov Decision Process (MDP) is defined by a tuple  $(\mathcal{S},\mathcal{A},\mathcal{P},\mathcal{R},\mu ,\gamma)$ , where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is the action space,  $\mathcal{P}:\mathcal{S}\times \mathcal{A}\to \mathcal{S}$  is the transition function,  $\mathcal{R}:S\times \mathcal{A}\times S\rightarrow \mathbb{R}$  is the reward function for any transition,  $\gamma \in (0,1]$  is a discount factor, and  $\mu$  is the initial state distribution. At each timestep  $t$ , the agent chooses an action  $a_{t}$  by following the policy  $\pi :S\to \Delta_{\mathcal{A}}$ . Then the agent obtains the next state  $s_{t + 1}$  and receives a scalar reward  $r_t$ . In single-task RL, the goal is to learn a policy  $\pi^{*} = \arg \max_{\pi}\mathbb{E}_{a_{t}\sim \pi}\left[\sum_{t = 0}^{\infty}\gamma^{t}r_{t}\right]$  by maximizing the expected cumulative reward of the corresponding task.
In a multi-task setting, different tasks can have different reward functions, state spaces and transition functions. We consider all tasks to share the same action space with the same embodied agent. Given a specific task  $\mathcal{T} \sim p(\mathcal{T})$ , a task-specified MDP can be defined as  $(\mathcal{S}^{\mathcal{T}}, \mathcal{A}, \mathcal{P}^{\mathcal{T}}, \mathcal{R}^{\mathcal{T}}, \mu^{\mathcal{T}}, \gamma)$ . Instead of solving a single MDP, the goal of multi-task RL is to find an optimal policy that maximizes expected return over all the tasks:  $\pi^{*} = \arg \max_{\pi} \mathbb{E}_{\mathcal{T} \sim p(\mathcal{T})} \mathbb{E}_{a_{t} \sim \pi^{\mathcal{T}}} \left[ \sum_{t=0}^{\infty} \gamma^{t} r_{t}^{\mathcal{T}} \right]$ .
Multi-Task Offline Decision-Making. In offline decision-making, the policy is learned from a static dataset of transitions  $\{(s_j,a_j,s_j',r_j)\}_{j = 1}^N$  collected by an unknown behavior policy  $\pi_{\beta}$  [30]. In the multi-task offline RL setting, the dataset  $\mathcal{D}$  is partitioned into per-task subsets as  $\mathcal{D} = \cup_{i = 1}^{N}\mathcal{D}_i$  where  $\mathcal{D}_i$  consists of experiences from task  $\mathcal{T}_i$ . The key issue of RL in the offline setting is the distribution shift problem caused by temporal-difference (TD) learning. In our work, we extend the idea of Decision Diffuser [2] by considering multi-task policy learning as a conditional generative process without fitting a value function. The insight is to take advantage of the powerful distribution modeling ability of diffusion models for multi-task data, avoiding facing the risk of distribution shift.
Offline RL learns policies from a static dataset, which makes the quality and diversity of the dataset significant [15]. One can perform data perturbation [28] to up-sample the offline dataset. Alternatively, we synthesize new transitions  $(s,a,s^{\prime},r)$  by capturing the underlying MDP of a given task via diffusion models, which expands the original dataset and leads to significant policy improvement.
We employ diffusion models to learn from multi-task data  $\mathcal{D} = \cup_{i=1}^{N}\mathcal{D}_{i}$  in this paper. With  $\tau$  the sampled trajectory from  $\mathcal{D}$ , we denote  $\pmb{x}_{k}(\tau)$  as the  $k$ -step denoised output of the diffusion model, and  $\pmb{y}(\tau)$  is the condition which represents the task attributes and the trajectory's optimality (e.g., returns). A forward diffusion chain gradually adds noise to the data  $\pmb{x}_{0}(\tau) \sim q(\pmb{x}(\tau))$  in  $K$  steps with a pre-defined variance schedule  $\beta_{k}$ , which can be expressed as
$$
q \left(\boldsymbol {x} _ {k} (\tau) \mid \boldsymbol {x} _ {k - 1} (\tau)\right) := \mathcal {N} \left(\boldsymbol {x} _ {k} (\tau); \sqrt {1 - \beta_ {k}} \boldsymbol {x} _ {k - 1} (\tau), \beta_ {k} \boldsymbol {I}\right). \tag {1}
$$
In this paper, we adopt Variance Preserving (VP) beta schedule [67] and define  $\beta_{k} = 1 - \exp\left(-\beta_{min}\left(\frac{1}{K}\right) - 0.5\left(\beta_{\max} - \beta_{\min}\right)\frac{2k - 1}{K^{2}}\right)$ , where  $\beta_{\max} = 10$  and  $\beta_{\min} = 0.1$  are constants. A trainable reverse diffusion chain, constructed as  $p_{\theta}(\boldsymbol{x}_{k - 1}(\tau)|\boldsymbol{x}_k(\tau),\boldsymbol{y}(\tau)) := \mathcal{N}(\boldsymbol{x}_{k - 1}(\tau)|\mu_{\theta}(\boldsymbol{x}_k(\tau),\boldsymbol{y}(\tau),k),\Sigma_k)$ , can be optimized by a simplified surrogate loss [20]:
$$
\mathcal {L} _ {\text {d e n o i s e}} := \mathbb {E} _ {k \sim \mathcal {U} (1, K), \boldsymbol {x} _ {0} (\tau) \sim q, \epsilon \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I})} [ \| \epsilon - \epsilon_ {\theta} (\boldsymbol {x} _ {k} (\tau), \boldsymbol {y} (\tau), k) \| ^ {2} ], \tag {2}
$$
where  $\epsilon_{\theta}$  parameterized by a deep neural network is trained to predict the noise  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{I})$  added to the dataset sample  $\pmb{x}_0(\tau)$  to produce  $\pmb{x}_k(\tau)$ . By setting  $\alpha_{k}\coloneqq 1 - \beta_{k}$  and  $\bar{\alpha}_{k}\coloneqq \prod_{s = 1}^{k}\alpha_{s}$ , we obtain
$$
\boldsymbol {x} _ {k - 1} (\tau) \leftarrow \frac {1}{\sqrt {\alpha_ {k}}} \left(\boldsymbol {x} _ {k} (\tau) - \frac {\beta_ {k}}{\sqrt {1 - \bar {\alpha} _ {k}}} \epsilon_ {\theta} (\boldsymbol {x} _ {k} (\tau), \boldsymbol {y} (\tau), k)\right) + \sqrt {\beta_ {k}} \sigma , \sigma \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}), \text {f o r} k = \{K, \dots , 1 \}.
$$
Classifier-free guidance [19] aims to learn the conditional distribution  $q(\pmb{x}(\tau) | \pmb{y}(\tau))$  without separately training a classifier. In the training stage, this method needs to learn both a conditional  $\epsilon_{\theta}(\pmb{x}_k(\tau), \pmb{y}(\tau), k)$  and an unconditional  $\epsilon_{\theta}(\pmb{x}_k(\tau), \emptyset, k)$  model, where  $\pmb{y}(\tau)$  is dropped out. Then the perturbed noise  $\epsilon_{\theta}(\pmb{x}_k(\tau), \emptyset, k) + \alpha (\epsilon_{\theta}(\pmb{x}_k(\tau), \pmb{y}(\tau), k) - \epsilon_{\theta}(\pmb{x}_k(\tau), \emptyset, k))$  is used to generate samples latter, where  $\alpha$  can be recognized as the guidance scale.
Figure 2: An example of the denoising process of MTDIFF. We choose 4 tasks for visualization and  $\pmb{x}_K(\tau)$  is sampled from the Gaussian noise for each task. Since different tasks require different manipulation skills, the corresponding action sequences are dispersed in the embedding space. Our model learns such properties and generates task-specific sequences based on task-relevant prompts.
3.1 Diffusion Formulation
To capture the multimodal distribution of the trajectories sampled from multiple MDPs, we formulate the multi-task trajectory modeling as a conditional generative problem via diffusion models:
$$
\max  _ {\theta} \mathbb {E} _ {\tau \sim \cup_ {i} \mathcal {D} _ {i}} \left[ \log p _ {\theta} \left(\boldsymbol {x} _ {0} (\tau) \mid \boldsymbol {y} (\tau) \right], \right. \tag {3}
$$
where  $\pmb{x}_0(\tau)$  is the generated desired sequence and  $\pmb{y}(\tau)$  is the condition.  $\pmb{x}_0(\tau)$  will then be used for generative planning or data synthesis through conditional reverse denoising process  $p_{\theta}$  for specific tasks. Maximizing Eq. (3) can be approximated by maximizing a variational lower bound [20].
In terms of different inputs and outputs in generative planning and data synthesis,  $\pmb{x}(\tau)$  can be represented in different formats. We consider two choices to formulate  $\pmb{x}(\tau)$  in MTDIFF-P and MTDIFF-S, respectively. (i) For MTDIFF-P,  $\pmb{x}(\tau)$  represents the action sequence for planning. We model the action sequence defined as:
$$
\boldsymbol {x} _ {k} ^ {p} (\tau) := \left(a _ {t}, a _ {t + 1}, \dots , a _ {t + H - 1}\right) _ {k}, \tag {4}
$$
with the context condition as
$$
\boldsymbol {y} ^ {p} (\tau) := \left[ \boldsymbol {y} ^ {\prime} (\tau), R (\tau) \right], \quad \boldsymbol {y} ^ {\prime} (\tau) := (Z, s _ {t - L + 1}, \dots , s _ {t}), \tag {5}
$$
where  $t$ ,  $H$ ,  $R(\tau)$  and  $L$  denote the time visited in trajectory  $\tau$ , the length of the input sequence  $\mathbf{x}$ , the normalized cumulative return under  $\tau$  and the length of the observed state history, respectively.  $Z$  is the task-relevant information as prompt. We use  $\mathbf{y}'(\tau)$  as an ordinary condition that is injected into the model during both training and testing, while considering  $R(\tau)$  as the classifier-free guidance to obtain the optimal action sequence for a given task. (ii) For data synthesis in MTDIFF-S, the inputs and outputs become the transition sequence that contains states, actions, and rewards, and then the outputs are utilized for data augmentation. We define the transition sequence as:
$$
\boldsymbol {x} _ {k} ^ {s} (\tau) := \left[ \begin{array}{c c c c} s _ {t} & s _ {t + 1} & \dots & s _ {t + H - 1} \\ a _ {t} & a _ {t + 1} & \dots & a _ {t + H - 1} \\ r _ {t} & r _ {t + 1} & \dots & r _ {t + H - 1} \end{array} \right], \tag {6}
$$
with the condition:
$$
\boldsymbol {y} ^ {\mathrm {s}} (\tau) := [ Z ], \tag {7}
$$
where  $y^{s}(\tau)$  takes the same conditional approach as  $y^{\prime}(\tau)$ . Figure 2 illustrates the reverse denoising process of MTDIFF-P learned on multi-task datasets collected in Meta-World [73]. The result demonstrates that our diffusion model successfully distinguishes different tasks and finally generates the desired  $x_0(\tau)$ . We illustrate the data distribution of  $x_0(\tau)$  in a 2D space with dimensional reduction via T-SNE [61], as well as the rendered states after executing the action sequence. The result shows that, with different task-specific prompts as conditions, the generated planning sequence for a specific task will be separate from sequences of other tasks, which verifies that MTDIFF can learn the distribution of multimodal trajectories based on  $y(\tau)$ .
3.2 Prompt, Training and Sampling
In multi-task RL and LLM-driven decision-making, existing works use one-hot task identifiers [56, 74] or language descriptions [1, 6] as conditions in multi-task training. Nevertheless, we argue that the one-hot encoding for each task [22, 14] suffices for learning a repertoire of training tasks while cannot generalize to novel tasks since it does not leverage semantic similarity between tasks. In addition, the language descriptions [51, 1, 6, 41, 40] of tasks require large amounts of human labor to annotate and encounter challenges related to ambiguity [72]. In MTDIFF, we use expert demonstrations consisting of a few trajectory segments to construct more expressive prompts in multi-task settings. The incorporation of prompt learning improves the model's ability for generalization and extracting task-relevant information to facilitate both generative planning and data synthesis. We remark that a similar method has also been used in PromptDT [69]. Nonetheless, how such a prompt contributes within a diffusion-based framework remains to be investigated.
Specifically, we formulate the task-specific label  $Z$  as trajectory prompts that contain states and actions:
$$
Z := \left[ \begin{array}{c c c c} s _ {i} ^ {*} & s _ {i + 1} ^ {*} & \dots & s _ {i + J - 1} ^ {*} \\ a _ {i} ^ {*} & a _ {i + 1} ^ {*} & \dots & a _ {i + J - 1} ^ {*} \end{array} \right], \tag {8}
$$
where each element with star-script is associated with a trajectory prompt, and  $J$  is the number of environment steps for identifying tasks. With the prompts as conditions, MTDIFF can specify the task by implicitly capturing the transition model and the reward function stored in the prompts for better generalization to unseen tasks without additional parameter-tuning.
In terms of decision-making in MTDIFF-P, we aim to devise the optimal behaviors that maximize return. Our approach is to utilize the diffusion model for action planning via classifier-free guidance [19]. Formally, an optimal action sequence  $\pmb{x}_0^p(\tau)$  is sampled by starting with Gaussian noise  $\pmb{x}_K(\tau)$  and refining  $\pmb{x}_k^p(\tau)$  into  $\pmb{x}_{k-1}^p(\tau)$  at each intermediate timestep with the perturbed noise:
$$
\epsilon_ {\theta} \left(\boldsymbol {x} _ {k} ^ {p} (\tau), \boldsymbol {y} ^ {\prime} (\tau), \varnothing , k\right) + \alpha \left(\epsilon_ {\theta} \left(\boldsymbol {x} _ {k} ^ {p} (\tau), \boldsymbol {y} ^ {\prime} (\tau), R (\tau), k\right) - \epsilon_ {\theta} \left(\boldsymbol {x} _ {k} ^ {p} (\tau), \boldsymbol {y} ^ {\prime} (\tau), \varnothing , k\right)\right), \tag {9}
$$
where  $\pmb{y}^{\prime}(\tau)$  is defined in Eq. (5).  $R(\tau)$  is the normalized return of  $\tau$ , and  $\alpha$  is a hyper-parameter that seeks to augment and extract the best portions of trajectories in the dataset with high return. During training, we follow DDPM [20] as well as classifier-free guidance [19] to train the reverse diffusion process  $p_{\theta}$ , parameterized through the noise model  $\epsilon_{\theta}$ , with the following loss:
$$
\mathcal {L} ^ {p} (\theta) := \mathbb {E} _ {k \sim \mathcal {U} (1, K), \boldsymbol {x} _ {0} (\tau) \sim q, \epsilon \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}), \beta \sim \operatorname {B e r n} (\mathrm {p})} \left[ \| \epsilon - \epsilon_ {\theta} \left(\boldsymbol {x} _ {k} ^ {p} (\tau), \boldsymbol {y} ^ {\prime} (\tau), (1 - \beta) R (\tau) + \beta \emptyset , k\right) \| ^ {2} \right]. \tag {10}
$$
Note that with probability  $p$  sampled from a Bernoulli distribution, we ignore the conditioning return  $R(\tau)$ . During inference, we adopt the low-temperature sampling technique [2] to produce high-likelihood sequences. We sample  $\pmb{x}_{k - 1}^{p}(\tau)\sim \mathcal{N}(\mu_{\theta}(\pmb{x}_{k - 1}^{p},\pmb{y}^{\prime}(\tau),R_{\max}(\tau),k - 1),\beta \Sigma_{k - 1})$  where the variance is reduced by  $\beta \in [0,1)$  for generating action sequences with higher optimality.
For MTDIFF-S, since the model aims to synthesize diverse trajectories for data augmentation, which does not need to take any guidance like  $R(\tau)$ , we have the following loss:
$$
\mathcal {L} ^ {s} (\theta) := \mathbb {E} _ {k \sim \mathcal {U} (1, K), \boldsymbol {x} _ {0} (\tau) \sim q, \epsilon \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I})} \left[ \left\| \epsilon - \epsilon_ {\theta} \left(\boldsymbol {x} _ {k} ^ {s} (\tau), \boldsymbol {y} ^ {s} (\tau), k\right) \right\| ^ {2} \right]. \tag {11}
$$
We sample  $\pmb{x}_{k - 1}^{s}(\tau)\sim \mathcal{N}(\mu_{\theta}(\pmb{x}_{k - 1}^{s},\pmb{y}^{s}(\tau),k - 1),\Sigma_{k - 1})$  . The evaluation process is given in Fig. 3.
3.3 Architecture Design
Notably, the emergence of Transformer [62] and its applications on generative modeling [44, 5, 47, 6] provides a promising solution to capture interactions between modalities of different tasks. Naturally, instead of U-Net [48] which is commonly used in previous single-task diffusion RL works [21, 2, 33], we parameterize  $\epsilon_{\theta}$  with a novel transformer architecture. We adopt GPT2 [45] architecture for implementation, which excels in sequential modeling and offers a favorable balance between performance and computational efficiency. Our key insight is to train the diffusion model in a unified manner to model multi-task data, treating different inputs as tokens in a unified architecture, which is expected to enhance the efficiency of diverse information exploitation.
Figure 3: Model architecture of MTDIFF, which treats different inputs as tokens in a unified architecture. The two key designs are (i) the trainable GPT2 Transformer which enhances sequential modeling, and (ii) the MLPs and prediction head which enable efficient training.
As shown in Figure 3, first, different raw inputs  $x$  are embedded into embeddings  $h$  of the same size  $d$  via separate MLPs  $f$ , which can be expressed as:
$h_P = f_P(x^{\mathrm{prompt}}), h_{Ti} = f_{Ti}(x^{\mathrm{timestep}})$ ,  $\triangleright$  for prompt and diffusion timestep
$h_{Tr}^{s} = f_{Tr}(x^{\mathrm{transitions}})$ $\triangleright$  for MTDIFF-S
$h_A^p = f_A(x^{\mathrm{actions}}), h_H^p = f_H(x^{\mathrm{history}}), h_R^p = f_R(x^{\mathrm{return}}). \triangleright$  for MTDIFF-P
Then, the embeddings  $h_P$  and  $h_{Ti}$  are prepended as follows to formulate input tokens for MTDIFF-P and MTDIFF-S, respectively:
$$
h _ {\text {t o k e n s}} ^ {p} = \mathrm {L N} \left(h _ {T i} \times \left[ h _ {P}, h _ {T i}, h _ {R} ^ {p}, h _ {H} ^ {p}, h _ {A} ^ {p} \right] + h _ {R} ^ {p} + E ^ {\text {p o s}}\right), h _ {\text {t o k e n s}} ^ {s} = \mathrm {L N} \left(h _ {T i} \times \left[ h _ {P}, h _ {T i}, h _ {T r} ^ {s} \right] + E ^ {\text {p o s}}\right),
$$