Graph World Model

Response to Issue 1 (Missing details): Thanks for the comments. The settings for the baselines primarily follow LLAGA (Chen et al., 2024a) and OFA (Liu et al., 2023a). As stated in Appendix A.3, we convert all nodes and labels in the Cora, PubMed, and HIV datasets into text. For all methods, we use BERT to obtain text embedding. We divide all datasets into training, validation, and test sets in an 8:1:1 ratio. We will add these details to the next version.

Response to Issue 2 (Experiment costs): Thanks for your questions. We answer them one by one:

[Training and inference efficiency.] Accurately comparing the training and inference efficiency of GWM with other FMs is highly challenging because many FMs are not designed to address multimodal problems and utilize various architectures. We can only compare the efficiency between GWM and LLM-based FMs from principles. For GWM-T, its efficiency shows no fundamental difference from other LLM-based FMs, as both are based on the standard instruction tuning. For GWM-E, its training process only requires fine-tuning the projector as stated in [lines 275-277], and its embedding-based method also saves a significant amount of token cost, making it more efficient. For GWM-E, it takes ~7 hours (~$\frac{1}{4}$ of GWM-T) of training time on four NVIDIA A6000 GPUs described in [lines 352-376], and the inference time per case averages 0.213s (similar to GWM-T). Moreover, GWM-E significantly reduces memory usage with a shorter token length of 140.23 (~$\frac{1}{14}$ of GWM-T).

[Large-scale graphs feasibility.] GWM-E can be deployed to large-scale graphs. In section 4.2, we designed a simplified GCN. It not only reduces computational complexity (Wu et al., 2019; He et al., 2020a) but also benefits from PyTorch's sparse matrix multiplication acceleration, allowing it to scale efficiently to large graphs. As shown in Tables 10 and 11, we have conducted experiments with large graphs at the scale of hundreds of thousands, demonstrating GWM's ability to address computational complexity and scalability.

Response to Issue 3 (Examples) and 4 (A mInor issue): Thanks for the feedback. These tasks can all be framed as generation tasks when using LLM to construct a GWM. But as stated in section 2.3, the essence of world optimization is to guide agents in decision-making, while world generation primarily enables agents to perform QA responses and generation, thus, their objectives are different. We distinguish these tasks to better demonstrate the capabilities of GWM in various aspects. Additionally, we have provided detailed descriptions of each task in section 2.3, section 5, Table 9, and Tables 13-22 of the paper. However, following the reviewers' suggestions, we have supplemented the details in Tables 1-7 of rebuttal.pdf. We will add these details to the next version.

Response to utilization of graph structures: Thanks for the comment. Both GWM versions fully leverage graph structures for message passing from the perspectives of tokens and embeddings. Furthermore, we verify through the experiment in section 5.2 (No Graph vs N-hop) that graph structures enhance performance on downstream tasks.

Response to Question 1: Thanks for your question. We didn't include INSTRUCTG2I (with complex QFormer design) as a baseline since our goal was to design a lightweight framework like GWM-E. Following your advice, we compared performance on multi-modal generation tasks as follows. Results show INSTRUCTG2I outperforms GWM on Goodreads but underperforms on Multi-Modal-Paper. This may be because the large-parameter INSTRUCTG2I also requires a larger dataset to achieve good results. In contrast, our lightweight GWM offers better efficiency while adapting well to multiple scenarios. We'll include this discussion in our next version.

Response to Question 2: To validate whether similar conclusions to those in [1] can be achieved on GWM, we conducted an ablation study. Specifically, we designed baselines named Text-only, which means training and testing solely on the text modality compared with GWM-E. We made comparisons in planning and optimization (ALFWorld) as follows. It can be observed that removing other multimodal information leads to a performance decline in GWM, which aligns with the conclusions of [1]. We will conduct more ablation studies in the next version to further substantiate this finding.

Response to Weaknesses 1: Thanks for your questions. We answer your questions one by one:

[Rationale of cross-modal fusion of GWM.] There are many ways to perform multimodal fusion. We selected two representative methods, not to avoid advanced fusion techniques. One is a simple and direct method that unifies multimodal information into text via GWM-T. The other method uses GWM-E, employing a cross-modal projector to unify different modal embeddings efficiently. Moreover, GWM-E does not have a modality dominance issue. We designed baselines named No-image and No-text, which mean training and testing solely on the text modality or image modality compared with GWM-E. We made comparisons on ALFWorld as follows. We can observe that the impact of images and text on GWM-E is almost equivalent, which demonstrates that GWM-E does not have modality dominance and all modalities are important. We will conduct more ablation studies in the next version to further substantiate this finding.

[Rationale of simplified GCN and its adaptiveness to heterogeneous graphs.] Indeed, our main experiment is based on homogeneous graphs. Because heterogeneous graphs often rely on predefined node and edge types, but we have multiple tasks and also need the model to generalize to unseen tasks, it is hard to define node and edge types in advance. If heterogeneous type information is represented as tokens or embeddings similar to GWM, ordinary GNNs may solve this, but it would lead to complexity in implementation and be costly. A potential way to extend directly from GWM-E to heterogeneous graphs is to perform separate multi-hop aggregations based on different types of edges and then flatten the resulting node embeddings into sequences to feed into the LLM decoder. However, we believe that our results on six real-world tasks validate GWM’s effectiveness and efficiency, highlighting our contribution. We will leave the exploration of heterogeneous graphs to future work.

Response to Weaknesses 2: Thanks for your comments. We have already detailed how GWM models state transitions in section 2.2. As described in [lines 56-73], we use a graph for general multimodal state modeling. Moreover, to address the diversity of action settings in traditional WM, we model actions as action nodes ([lines 75-88]) and unify them into text descriptions as in [section 3.3 and section 4.3]. This allows GWM to use a unified model to solve multi-tasks. Lastly, in [lines 89-93 of section 2.2], we explain that our GWM essentially learns a transition function, which is aligned with principled world models.

Response to Weaknesses 3: Thanks for your comments. We answer your questions step by step:

[Avoid over-smoothing.] As introduced in [lines 247-253] of section 4.2, we input the multi-hop embeddings into the decoder. This is equivalent to a skip connection, a standard practice that alleviates oversmoothing [1, 2].

[GWM handles non-homophilous graphs.] Essentially, both versions of GWM transform graphs into a standard form interpretable by LLMs. Existing work (Liu et al., 2023a; Chen et al., 2024a) has demonstrated that LLMs can also effectively understand graphs, thus we anticipate that LLMs may outperform GNNs in comprehending homophilous or non-homophilous graphs. But if we must understand homophilous graphs from the perspective of GNNs, a potential way is to integrate [3] based on GWM-E, but this would introduce complex implementation and significant costs. Moreover, the performance across six representative real-world tasks has already validated the practicability and generalizability of GWM. We will include the above discussions in the next version.

[Graph transitions align with task objectives.] As discussed in section 2.3, our action is essentially a description of the task objective. For example, in the RAG task, the action of GWM is to answer specific queries with retrieved contexts. Thus, the graph transition modelled by GWM is aligned with the task objective without a theoretical gap.

[1] Optimization of graph neural networks: Implicit acceleration by skip connections and more depth.

[2] Representation learning on graphs with jumping knowledge networks.

[3] Revisiting heterophily for graph neural networks.

Response to missing theory and hyperparameters analysis: Thanks for the comment. We prioritize empirical analysis of GWM across multiple tasks over theoretical proofs, which are challenging in LLM research. Given the foundation model's numerous hyperparameters and computational costs, we focused on introducing key hyperparameters (lines 357-372 and Appendix B) and analyzed critical ones like hop num (section 5.2), balancing reproducibility with practical costs.

Q1. Can the author supplement the baseline that tuned on the task-specific dataset (e.g., Table 6).

Response: Thanks for the comments. Indeed, we have detailed the specific settings of the baselines for Table 6 in Appendix A.4 ([lines 779-796]). We primarily selected three classic agent baselines for comparison: CoT, ToT, and Few-shots. Although these baselines were not additionally fine-tuned on this dataset, the results of GWM shown in Figure 6, as described in [lines 382-384], were also not additionally fine-tuned. Therefore, from this perspective, the comparison between GWM and the baselines is fair. However, following the reviewer's suggestion, we compared two baselines fine-tuned on AgentClinic: FT is a baseline finetuned on the task using the Llama-3-8B model. Longformer [1] is a strong baseline for long document understanding. The table below reports the comparison results between GWM and these baselines, demonstrating the effectiveness of GWM in AgentClinic.

[1] Longformer: The long-document transformer.

Q2. The authors should have included baseline graph world models (Zhang et al. 2021; Zhu et al. ICLR 2023) in their experiments.

Response: Thanks for your feedback. Indeed, the works on graph world models (Zhang et al. 2021; Zhu et al. ICLR 2023) do not consider the processing of multimodal information in their modeling and primarily focus on offline scenarios in the RL field. These limitations restrict their capabilities as world models and also make them inapplicable for direct comparison in our scenario. In contrast, GWM is a world model capable of handling multimodal information and generalizing across various tasks. We will supplement these discussions in the related work section of the next version.

Q3. The paper fails to address computational complexity and scalability considerations adequately.

Response: Thanks for your valuable comments. In fact, we address computational complexity and scalability from two perspectives. First, we introduce GWM-E, which, as discussed in [lines 275-277], fixes the LLM's parameters and only fine-tunes the projector parameters, thereby reducing computational complexity. Additionally, as described in section 4.2 [lines 234-246], we designed a simplified, parameter-free GCN. This not only reduces computational complexity compared to traditional GNNs (Wu et al., 2019; He et al., 2020a) but also benefits from PyTorch's sparse matrix multiplication acceleration, allowing it to scale more effectively to large graphs. As shown in Tables 10 and 11 in the appendix, we have conducted experiments with large graphs at the scale of hundreds of thousands, demonstrating GWM's ability to address computational complexity and scalability.

Q4. The authors present only quantitative results without providing qualitative comparisons.

Response: Thanks for the insightful advice. We have already added qualitative comparisons for all tasks in Tables 1-7 of rebuttal.pdf. We will complete this information in the appendix of the next version.

Q5. Figure 3, which illustrates the GWM framework, suffers from quality issues that impede understanding. (1) Target nodes are not clearly specified—variables should be directly labeled. Additionally, the target nodes and prompt in the top-right of Figure 3 are two separate inputs but are not represented with distinct arrows. (2) The use of Stable Diffusion in the top-right creates confusion since it was designed for text-to-image generation, while the diagram shows token inputs and "next state" outputs. The authors should clarify that the model operates in Stable Diffusion's latent space.

Response: Thanks for your valuable suggestions. Following your suggestions, we first specified the target nodes in Figure 3. Next, regarding the issue with the top right arrow, we intended to prompt the target nodes, so we changed the content on the arrow to 'Prompt (Target nodes)'. Lastly, concerning the input and output issues of Stable Diffusion in GWM-T, as described in [lines 229-234], our input included action nodes and target nodes at the token level, and the output image served as the next state. To further clarify this process, we added legends to distinguish between different colors and shapes. We have updated the modified figure in Figure 2 of rebuttal.pdf and it will be updated in the next version.

Q1. The method, although different in nature, solves a similar task to LANISTR [a], which would be nice to cite and contrast against.

Response: Thanks for the valuable suggestions. Indeed, we have already compared two baselines that, like LANISTR, were pre-trained by modality alignment and then fine-tuned on downstream tasks, Contrastive MLP (Liu et al., 2022), CLIP FT (Radford et al., 2021) in [lines 348-352] and [Table 3]. However, we followed the reviewer's suggestions and compared the performance of GWM with LANISTR. Due to the limited time for rebuttal, we primarily focused on comparing the performance of GWM and LANISTR when directly finetuned on the Goodreads dataset in the task of multi-modal matching. It can be observed that GWM achieved a significant performance gain compared to LANISTR. In the next version, we will compare LANISTR in more tasks and evaluate its performance after pretraining on our dataset.

Q2. In Tab. 4, the provided baselines are old. Maybe the authors can include FREEDOM [e] or another recent SOTA model for recommendation.

Response: Thank you for your insightful comments. In the recommendation task, we primarily selected some classic and representative baselines for comparison. Our goal in comparison with task-specific baselines is not to demonstrate that GWM can beat sota in all tasks, but to show that GWM, with a unified model, has good generalization capabilities across many tasks, including recommendation. However, we followed the reviewer's suggestions and compared the performance of GWM with FREEDOM on the recommendation task shown in the following table. We can observe that, compared to FREEDOM, the two methods of GWM achieved better performance, which demonstrates the effectiveness of GWM. We will update the experimental results and discussion in the next version.

Q3. The graphics can be improved to make the method clearer and more understandable. Eg. in Figs. 1,3 it is not clear what each color identifies in Fig. 1: what is the difference between intended and unintended action, and what is the difference between the node/edge/graph level actions?

Response: Thanks for your valuable advice. Following your suggestions, we have separated the descriptions of different action levels and the methods for obtaining target nodes in Figure 1. Additionally, we have standardized the colors in corresponding parts of Figures 1 and 3. Lastly, we have added captions to each figure to differentiate the meanings of various colors and shapes. We showed the figures in Figure 1 and 2 of rebuttal.pdf. We will update these modifications in the next version.