Decision Mixer: Integrating Long-term and Local Dependencies via Dynamic Token Selection for Decision-Making

Thanks for the careful review of our work. Due to the strict word limit, we have tried to address the reviewers' comments carefully. All additional experiments will be incorporated into the text.

Theoretical Claims
We provided a brief analysis of the CSM method from the perspective of re-weighting in equation (2). Given the word limit and the need for academic rigor, we have removed the statement "theoretical analysis" from the main text.

Concern 1: The inconsistency of training and inference
We carefully considered using a more consistent approach to achieve our objectives. We designed three schemes: (1) The hypernetwork uses a single token as input for training and inference without needing an auxiliary predictor. (2) The hypernetwork uses a causal sequence as input for training and inference without needing an auxiliary predictor. (3) The hypernetwork uses the entire sequence as input for training, providing data for the binary classification training of the auxiliary predictor, which is then used for inference (adopted by DM).

(3) demonstrated the best performance and stability across all tasks, with significantly shorter training times. Although (1) and (2) ensured consistency between training and inference, they struggled with convergence due to insufficient utilization of global information from the training data. (3) adopts a task decomposition approach to implement training and inference hierarchically, with the auxiliary predictor trained based on the prediction data from each round of the hypernetwork and router. Figure 1 of the pdf (https://anonymous.4open.science/r/Decision-Mixer-1068/rebuttal.pdf) shows the training loss curves, which indirectly suggest that binary classification of data on a specific Mixer layer for a given task is relatively easy to learn. Experimental results in Table 1 and the ablation study in Table 2 confirm that potential distribution shifts were addressed through synchronized training.

Concern 2: The novelty and contribution seem to be incremental
All prior works have tended to introduce expert routing mechanisms in FFN or attention layers without involving token selection, where a fixed number of experts handle different tokens in the complete sequence. We emphasize that DM differs fundamentally from existing works in component design and training approach. DM innovatively designs (1) a dynamic token selection mechanism to address sequence modeling conflicts specific to offline RL, differing from conventional static MoE. DM handles incomplete sequences and uses generalized residual connections after each layer to ensure the consistency of output and input lengths. (2) During inference, we also designed a unique auxiliary predictor from the task decomposition perspective to address inconsistencies between training and inference. (3) We deploy the selection before the attention layer, making DM a more flexible plug-and-play architecture. It is more cost-efficient than conventional combinations of MoE and transformers, providing a reasonable direction for exploring scaling laws under the DT architecture.

Essential References Not Discussed
We will discuss the combination of MoE and transformers in the text. GShard[1] introduced MoE into transformers to address load imbalance via routing and expert capacity constraints. Switch Transformers[2] proposed a Top-1 gating mechanism to reduce computation and communication overhead. DeepSeekMoE[3] optimized expert utilization with fine-grained segmentation and expert isolation. MoEUT[4] combined MoE with the Universal Transformer, addressing the parameter-computation efficiency trade-off. Additionally, several works[5,6,7] explored MoE integration in visual tasks from the perspectives of data processing[5], multi-task learning[6], pre-training[7], and global training[8].

We discuss MoE because DM can be viewed as using a single expert to filter tokens before the standard transformer layer. However, this does not mean we have merely made a simple transplant. The collaboration between the router and the hypernetwork in DM, along with the auxiliary predictor and generalized residual connections, are seamlessly integrated, forming an efficient and practical framework that paves the way for future paradigms in offline RL.

[1] Gshard...
[2] Switch transformers...
[3] Deepseekmoe...
[4] Moeut...
[5] Scaling vision with sparse mixture...
[6] Adamv-moe: Adaptive multi-task vision...
[7] Moe jetpack: From dense checkpoints...
[8] Mod-squad: Designing mixtures...

Thanks for the detailed review of our work. Given the strict word limit, we have carefully addressed the reviewer's comments and will incorporate the additional experiments and suggestions into the updated version. Anonymous pdf: https://anonymous.4open.science/r/Decision-Mixer-1068/rebuttal.pdf

The details of $R$ and $H$
We found that both $R$ and $H$ can perform well using simple MLP, as shown below. $R$ and $H$ are incorporated as part of the main model, with no additional constraints added, and the training is based solely on equation (6). Additional thoughts can be found in Q1 for Reviewer ekS2.

The auxiliary predictor $\theta^l_{aux}$ is trained with gradient isolation using dynamic selection results from $R$ and $H$. The main model selects the top-$k$ tokens from $X^l$ and generates binary labels $z_i \in {0,1}$ as ground truth for $\theta^l_{aux}$. The predictor outputs binary logits $\hat{y}_i$, optimized via equation (5) to match $\sigma(\hat{y}_i)$ with $z_i$. We found the token selection distribution in binary classification easy to learn, with rapid loss convergence shown in Figure 1 of the anonymous pdf. Main experiments and ablation studies confirm prediction accuracy. Alternatives are discussed in our response to reviewer EJaz on "Inconsistency of training and inference."

Lacks evaluation on specific environments
We have added experimental results for DM in the Atari, reporting the average performance across three random seeds. Table 1 of the anonymous pdf shows that DM performs well in discrete action environments, demonstrating the method's generalizability. This success was expected due to the precedents set by DT and DC.

More engage with the existing literature
We will provide a more detailed discussion in the main text. Token dropping, initially proposed to reduce BERT inference costs [1,2], was adapted by Hou et al. [3] to improve training efficiency. Random-LTD [4] advanced this with random-layer token dropping and learning rate scheduling. While prior work focuses on static efficiency strategies for vision and language tasks [5,6], they often lack dynamic adaptation, risking semantic disruption. In contrast, DM dynamically selects tokens using a router and hypernetwork, aligning with offline RL’s Markovian nature for better performance. Its plug-and-play design and synchronized training offer greater flexibility than existing methods.

[1] Train short, test long...
[2] SpAtten: Efficient Sparse...
[3] Token dropping for...
[4] Random-ltd...
[5] Revisiting token dropping...
[6] Multi-Stage Vision Token...

Further empirical evidence
In the paper, we have included baseline comparisons in RL scenarios, component ablation studies, token selection statistics and visualizations, computational efficiency, training curves, generalization, portability, and context length experiments. Following reviewer feedback, we added DM’s Atari results and adapted Random-LTD to DT in Gym. Table 2 of the anonymous pdf shows that Random-LTD underperforms DM in RL tasks, likely due to its reliance on random token dropping. We will attempt more transplant comparison experiments in future work.

Essential References Not Discussed
We will include LSDT in the related work and experimental comparison sections. LSDT integrates DT's self-attention and DC's dynamic convolution via branch design for decision-making. DM extends these environments and outperforms LSDT in most experiments (Table 3 of the anonymous pdf).

Other Comments

• We apologize for the typo and will change "MOD" to "model".
• We have placed the updated Figure 2 in the anonymous PDF.
• We will replace the relevant terms with \texit{k}.

Q1: Why does an auxiliary predictor improve performance
The auxiliary predictor addresses issues from inconsistent input data formats during training and inference. $H$ uses the complete sequence to predict $k$, but during inference, the autoregressive nature of the transformer makes the following sequence invisible, preventing sequence-level predictions. The token-level auxiliary predictor enables filtering without full sequence information. $R$ and $H$ jointly serve as teacher models, providing binary classification training data for the auxiliary predictor. This hierarchical training approach reduces the difficulty of training and fully utilizes all available data.

Q2: Include a clock-time measurement
We measured the clock time from training start to convergence across multiple tasks. The results in Table 4 (anonymous pdf) show that DM has a smaller time overhead than DT and is competitive with DC. Despite adding networks, DM's minimal parameters and dynamic token selection mechanism shorten the sequence length, avoiding significant clock time increases.

Thanks for the careful review of our work. Due to the strict word limit, we have tried to address the reviewers' comments carefully. All additional experiments and suggestions will be incorporated into the updated text.

It is recommended to adjust "MOE" to "MoE"
We apologize for this mistake. To ensure consistency and accuracy of terminology, we will change all occurrences of "MOE" to "MoE" in the main text.

Q1: Should regularization terms be introduced
We appreciate the reviewer's thoughtful insights. To ensure our algorithm's simplicity and ease of deployment, we have solely used Equation (6) for the entire model training. Introducing additional regularization terms, such as a routing weight smoothing term, can improve training stability. To investigate this, we conducted experiments by adding a regularization term with a weight of 0.1 on top of Equation (6). The specific details of the regularization are as follows: $L_{smooth}=\lambda \sum_{i=1}^{S-1}\|w_i-w_{i+1}\|^2$ The experimental results are presented in the table below. We refer to DM with the added regularization term as DM_s.

We found that while DM_s exhibited improved stability, it experienced varying degrees of performance degradation across all tasks compared to DM. We hypothesize that this is primarily due to the introduction of the weight smoothing term, which constrains the flexibility of token selection based on task-specific characteristics, making the selection process overly conservative. Abrupt changes in token selection across neighboring positions can be reasonable, as they allow the model to flexibly concatenate trajectories based on data characteristics while preserving the Markov properties required for specific tasks. In future work, we will explore more refined approaches to achieving efficient token selection more smoothly.

Q2: Is the dynamic range of k constrained
The range of values for $k$ is unrestricted to ensure the algorithm's simplicity. Depending on the nature of the task, $k$ can approach either 0 or the entire sequence length $S$. The medium-replay dataset (high noise), medium dataset (moderate noise), and medium-expert dataset (low noise) from specific tasks serve as references for evaluating the robustness of our model's selection, as shown in Table 8 and Figure 6. For the hopper-medium task, the average number of selected tokens fluctuates slightly across datasets of different qualities. A similar pattern is observed in the other two tasks, with the overall average number of selected tokens remaining between 20 and 30, approximately half of the entire sequence length, without significant degradation.

To further investigate the hypernetwork’s predictions in highly noisy environments, we introduced Gaussian noise sampled from a standard normal distribution to all tokens in the hopper-medium-replay and halfcheetah-medium-replay datasets while keeping the original action labels unchanged. We refer to these modified datasets as hopper-medium-noise and halfcheetah-medium-noise. The results show that the hypernetwork outputs a higher $k$ value for hopper-medium-noise, whereas the opposite trend is observed for halfcheetah-medium-noise. This suggests that the number of selected tokens adapts to task-specific characteristics in highly noisy environments. No significant performance degradation was observed during the experiments, indicating the robustness and stability of our approach.

W1/Q3: A more balanced discussion would strengthen the paper
We have summarized DM's advantages and limitations to help readers understand our approach better. By dynamically selecting and concatenating important tokens at each layer, DM reduces computational complexity while effectively balancing the trade-off between capturing long-term dependencies and extracting local Markov features. This approach enhances efficiency and offers valuable insights into scaling laws for offline RL. Notably, the dynamic token selection during inference relies on the auxiliary predictor for online decision-making, which may introduce latency in scenarios with strict real-time requirements. Additionally, DM performs slightly worse than value-based methods on low-quality or noisy data. Future work will explore data augmentation or more efficient and robust token selection strategies—such as adversarial training or noise-adaptive mechanisms—to improve adaptability to noisy or low-quality data.

Thanks for the careful review of our work. Due to the strict word limit, we have tried to address the reviewers' comments carefully. These additional experiments and suggestions will be incorporated into the updated main text.

W1: The robustness of the dynamic mechanism in sparse scenarios needs to be enhanced
We fully agree that there is still potential for improvement. An intuitive approach is to adopt data augmentation. Considering that complex tasks exhibit significant noise variations mainly at the environmental state level, we introduce Gaussian noise perturbation to all state dimensions in the training data, formulated as: $\tilde{s}_j=s_j+\alpha \cdot \sigma_j \cdot \epsilon_j.$ Here, $\alpha$ is a controllable noise intensity, which we set to 0.05. $\sigma_j$ represents the standard deviation of the $j$-th dimension, which has been computed during data preprocessing. $\epsilon \sim \mathcal{N}(0, I)$ denotes a noise vector sampled from the standard normal distribution. We generate two random seeds to add noise to the initial data $\tau$, obtaining two noisy datasets, $\hat{\tau_1}$ and $\hat{\tau_2}$. Each training iteration uses a total of $3 \times$ batch data. $\tau$, $\hat{\tau_1}$, and $\hat{\tau_2}$ share identical values except for the perturbed states. The processed data is then used for training and referred to as DM_enhanced. The experimental results on three tasks are as follows:

DM_enhanced performs better than DM on the maze2d-umaze task, with a minor standard deviation. Although DM_enhanced does not perform better on the reward-dense task hopper-medium, a similar stability improvement is observed. This phenomenon suggests that the data augmentation employed in DM_enhanced enhances its stability and robustness. Due to rebuttal time constraints, exploring more instructive designs for robustness suppression will be part of our future work.

Q1: Could the auxiliary predictor make incorrect selections due to distribution shift

The only potential distribution shift is the difference between the training and inference inputs. We designed three schemes: (1) The hypernetwork uses a single token as input for training and inference without needing an auxiliary predictor. (2) The hypernetwork uses a causal sequence as input for training and inference without needing an auxiliary predictor. (3) The hypernetwork uses the entire sequence as input for training, providing data for the binary classification training of the auxiliary predictor, which is then used for inference (the approach adopted by DM).

Scheme (3) demonstrated the best performance and stability across all tasks, with significantly shorter training times. Although schemes (1) and (2) ensured consistency between training and inference, they struggled with convergence due to insufficient utilization of global information from the training data. Scheme (3) adopts a task decomposition approach to implement training and inference hierarchically, with the auxiliary predictor trained based on the prediction data from each round of the hypernetwork and router. Figure 1 of the anonymous pdf shows the training loss curves, which indirectly suggest that binary classification of data on a specific Mixer layer for a given task is relatively easy to learn. Experimental results in Table 1 confirm that potential distribution shifts were addressed through synchronized training.

Q2: How does the "dynamic token selection mechanism" reduce computational overhead
In a specific Mixer Layer $l$, the dynamic token selection mechanism reduces computational complexity by adaptively selecting the top-$k$ tokens (via router $R^l$ and hypernetwork $H^l$) for attention processing while skipping others. When the sequence length input to the attention layer is reduced from $S$ to $k$, the attention complexity decreases from $O(S^2d)$ to $O(k^2d)$, and the computational complexity of the FFN layer reduces from $O(Sd)$ to $O(kd)$, where $d$ is the hidden dimension. For example, in the Key/Value projection stage, if $k = S/2$, the computation is reduced to 25% of the original FLOPs. Experiments show that DM reduces FLOPs by 47.0% compared to DT and achieves better clock time performance, as shown in the table below.