GL-Fusion: Rethinking GNN-LLM Integration

# 标准Llama-3: 严格下三角mask
[[1, 0, 0, 0],
 [1, 1, 0, 0],
 [1, 1, 1, 0],
 [1, 1, 1, 1]]

# GL-Fusion: 图内全连接 + 文本因果
[[1, 0, 0, ... | 0, 0, ... | 0],  ← 文本token
 [1, 1, 0, ... | 0, 0, ... | 0],
 [1, 1, 1, 1,  | 1, 1, ... | 0],  ← 图节点（内部全连接）
 [1, 1, 1, 1,  | 1, 1, ... | 1]]  ← 后续文本（看全部）

# Stage 1: 聚合节点内部256个tokens
query = hidden[token_i]  # [4096]
keys = node_text[node_j]  # [256, 4096]
weights = softmax(query @ keys.T)  # [256]
node_repr = weights @ keys  # [4096]

# Stage 2: 聚合50个节点
all_nodes = [node_repr_0, ..., node_repr_49]  # [50, 4096]
weights2 = softmax(query @ all_nodes.T)  # [50]
output = weights2 @ all_nodes  # [4096]

# GNN分支（位置20-69的图节点）
gnn_logits = Linear(hidden[target_pos])  # [num_classes]
# 优势：并行预测所有节点，适合分类任务

# LLM分支（位置-1的最后token）
lm_output = AutoRegressive_Generate(hidden[-1])
# 优势：生成自然语言，适合文本任务

# Ensemble
final_pred = combine(gnn_logits, lm_output)

Pipeline 1: 主序列（Llama-3 Tokenizer）
"Classify paper <graph_start><node>×50<graph_end> Answer:"
→ [73, 4096]  进入Transformer主干

Pipeline 2: 节点文本（LLM2Vec）
["Attention Is All You Need. The dominant...", ...]
→ [50, 256, 4096]  侧路，Cross-Attention时使用

Pipeline 3: 边特征（Edge Encoder）
["cites", "cites", "collaborated", ...]
→ [120, 4096]  侧路，Message Passing时使用

# Layer i (i=0..31)
hidden = structure_aware_self_attention(hidden)  # 全序列[73,4096]

if i in [0,4,8,12,16,20,24,28]:
    graph_part = hidden[:, 20:70, :]
    graph_part = message_passing(graph_part, edge_embeds)  # 用Pipeline 3
    hidden[:, 20:70, :] = graph_part  # 覆盖更新

if i in [3,7,11,15,19,23,27,31]:
    for pos in range(73):
        cross_out = cross_attention(hidden[pos], node_texts)  # 用Pipeline 2
        hidden[pos] += cross_out

hidden = ffn(hidden)

paper_12345 = {
    "title": "Attention Is All You Need",
    "abstract": "The dominant sequence transduction...",
    "label": 5,  # cs.CL (Computation and Language)
    "citations": [2341, 5678, 8901],  # 2跳采样→50节点子图
    "year": 2017
}

# 动态子图采样
subgraph = sample_k_hop_neighbors(target=12345, k=2, max_nodes=500)
# → 50个节点，120条边

# 三条Pipeline的输入
input_data = {
    "main_sequence": [
        "Classify paper into 40 categories. Graph:",
        "<graph_start>", "<node>"×50, "<graph_end>", "Answer:"
    ],  # [73] tokens

    "node_texts": [
        "Attention Is All You Need. The dominant...",  # 256 tokens
        "[Label: cs.CL] Neural Machine Translation...",  # 训练集加标签
        ...
    ],  # [50, 256]

    "edge_data": [
        (0, 1, "cites"), (0, 2, "cites"), ...
    ]  # 120条边
}

# 超参数
config = {
    "base_model": "Llama-3-8B",
    "lora_rank": 64,
    "learning_rate": 3e-5,
    "batch_size": 1,  # 子图大小不同
    "epochs": 1,  # 单epoch足够
    "optimizer": "AdamW",
    "weight_decay": 0.1
}

# 训练循环
for target_node in train_set:  # 90,941个样本
    # 1. 采样子图
    subgraph = sample_subgraph(target_node)

    # 2. 三路编码
    main_embeds = llama3_embedding(main_seq)
    node_embeds = llm2vec_encoder(node_texts)
    edge_embeds = edge_encoder(edges)

    # 3. 前向传播（32层Transformer）
    hidden = model(main_embeds, node_embeds, edge_embeds)

    # 4. 双损失
    gnn_loss = CE(gnn_classifier(hidden[20]), label)
    lm_loss = CE(lm_head(hidden[-1]), "Computation and Language")
    total_loss = gnn_loss + lm_loss

    # 5. 反向传播（只更新可训练部分）
    total_loss.backward()
    optimizer.step()

# 输入: "The cat sat"
tokens = [CLS, The, cat, sat]

# 标准因果mask
mask = [[1, 0, 0, 0],
        [1, 1, 0, 0],  # "cat"不能看到"sat"
        [1, 1, 1, 0],
        [1, 1, 1, 1]]

# Attention计算
scores = Q @ K.T / √d  # [4, 4]
masked = scores + (1-mask)*(-1e9)  # 上三角→-∞
weights = softmax(masked)  # [4, 4]，上三角→0
output = weights @ V  # [4, dim]

# Layer 0完整流程
def layer_0(hidden):
    # 时刻1: Self-Attention（全部[73,4096]参与）
    hidden = self_attention(hidden)

    # 时刻2: Message Passing（仅20:70参与）
    for node_u in range(50):
        neighbors = [node_v for v in adjacency[u]]
        msgs = [hidden[20+v] * edge_embeds[u,v] for v in neighbors]
        agg = concat([mean(msgs), max(msgs), std(msgs)])  # [12288]
        h_new = Linear(concat([hidden[20+u], agg]))  # [4096]
        hidden[20+u] = hidden[20+u] + gate*h_new  # 门控更新

    # 时刻3: FFN（全部参与）
    hidden = ffn(hidden)

    return hidden

# 可学习的投影矩阵
W_Q1 = nn.Linear(4096, 4096)  # 投影query（主序列）
W_K1 = nn.Linear(4096, 4096)  # 投影key（节点文本）

# 计算过程
query = hidden[70, :]  # "Answer:"的表示 [4096]
keys = node_texts[0, :, :]  # Node 0的256个tokens [256, 4096]

Q = W_Q1(query)  # [4096]
K = W_K1(keys)   # [256, 4096]

scores = Q @ K.T / √d  # [256]
weights = softmax(scores)  # 关注"Attention", "Transformer"等关键词
output = weights @ keys  # [4096]

# 梯度流分离
gnn_loss → gnn_classifier → message_passing → graph_nodes[20:70]
                                            ↘ main_transformer_lora

lm_loss → lm_head → text_tokens[70:73] → cross_attention → llm2vec
                                       ↘ main_transformer_lora

# 只在main_transformer_lora处汇合，梯度叠加但不冲突

# 错误理解
hidden[20:70] = hidden[20:70] + mp_output  # ✗

# 正确理解
hidden[20:70] = mp_output  # ✓ 覆盖更新

# 但MP内部有软残差
h_final = h_original + gate * h_new  # gate初始≈0