基于注意力机制与强化学习的TSP求解：从理论到工程实践

1次阅读

没有评论

共计 2656 个字符，预计需要花费 7 分钟才能阅读完成。

旅行商问题（TSP）是经典的组合优化问题，传统解法在大规模场景下面临严峻挑战：

动态规划 ：虽然能保证最优解，但时间复杂度为 O(n^2*2^n)，当节点数超过 20 时计算成本呈指数级增长
启发式算法 ：如蚁群算法、遗传算法等依赖大量迭代，且参数敏感，不同实例需要重新调参
数学规划 ：商用求解器（如 Gurobi）对 50+ 节点问题求解耗时可能超过小时级

这些方法普遍存在两个共性问题：无法利用历史求解经验，以及难以泛化到不同规模的问题实例。

深度学习为路由问题提供了新的解决思路，主流架构各有特点：

Pointer Networks：
优势：首个将注意力机制应用于组合优化的模型
劣势：固定长度的 RNN 编码器难以捕捉复杂城市分布关系
Transformer：
优势：全局注意力机制适合捕捉城市间长程依赖
劣势：计算复杂度随序列长度平方增长
图神经网络 (GNN)：
优势：显式建模城市间的拓扑结构
劣势：需要额外的图构建过程

我们的方案选择 Transformer 作为基础架构，因其在序列建模和并行计算方面的优势。

class TSPEncoder(nn.Module):
    def __init__(self, embed_dim=128, n_heads=8):
        super().__init__()
        self.city_embedding = nn.Linear(2, embed_dim)  # 输入坐标 (x,y)
        self.pos_encoding = PositionalEncoding(embed_dim)
        self.transformer = nn.TransformerEncoder(nn.TransformerEncoderLayer(embed_dim, n_heads),
            num_layers=3
        )

    def forward(self, cities):
        # cities: [batch_size, seq_len, 2]
        h = self.pos_encoding(self.city_embedding(cities))
        return self.transformer(h)  # [batch_size, seq_len, embed_dim]

关键点：
– 将二维坐标映射到高维空间
– 加入位置编码保留空间信息
– 使用多层 Transformer 提取城市间关系

class TSPDecoder(nn.Module):
    def __init__(self, embed_dim=128):
        super().__init__()
        self.context = nn.Parameter(torch.randn(embed_dim))
        self.gru = nn.GRUCell(embed_dim, embed_dim)
        self.attention = AdditiveAttention(embed_dim)

    def forward(self, enc_out, mask, prev_node=None):
        # enc_out: [batch_size, seq_len, embed_dim]
        if prev_node is None:
            h = self.context.expand(enc_out.size(0), -1)
        else:
            h = self.gru(enc_out[:, prev_node], h)

        # 计算注意力分布
        logits = self.attention(h.unsqueeze(1), enc_out, mask)
        return logits.squeeze(1)  # [batch_size, seq_len]

关键点：
– 使用 GRU 维持解码状态记忆
– 注意力机制动态选择下一个访问城市
– 掩码防止重复访问

采用 REINFORCE 算法配合基线技巧：

奖励设计 ：
使用路径长度的倒数作为奖励：reward = 1.0 / tour_length
对 batch 内奖励做归一化处理

损失函数 ：

def compute_loss(logits, actions, rewards):
    # logits: [batch_size, seq_len]
    # actions: [batch_size, seq_len]
    log_probs = -F.cross_entropy(logits, actions, reduction='none')
    advantage = rewards - rewards.mean()
    return -(log_probs * advantage).mean()

训练技巧 ：
采用课程学习（Curriculum Learning），从 10 节点开始逐步增加难度
使用熵正则化鼓励探索
学习率预热配合余弦退火

梯度检查点 ：在反向传播时重新计算中间结果

from torch.utils.checkpoint import checkpoint

def forward(self, x):
    return checkpoint(self._forward, x)

混合精度训练 ：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    logits = model(inputs)
    loss = compute_loss(logits, actions, rewards)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

在训练时采用不同规模实例（20-100 节点）混合训练，测试结果：