机器学习破解高阶魔方:扩散距离与束搜索的协同求解
1. 项目概述当机器学习“学会”看穿魔方的迷宫魔方这个诞生于上世纪的小小立方体至今仍让无数人着迷也让计算机科学家们头疼不已。从3×3的标准魔方到更庞大的4×4、5×5版本其状态空间以指数级膨胀形成了一个天文数字级别的图——凯莱图。在这个图中每一个可能的魔方状态是一个节点每一次转动便是在节点间移动。求解魔方本质上就是在这个庞大到无法遍历的图中寻找从任意一个打乱状态回到“复原态”的最短路径。传统的搜索算法如双向广度优先搜索在面对4×4魔方约7.4×10^45种状态时已力不从心更遑论5×5魔方约1.2×10^74种状态。过去几年以DeepCubeA和EfficientCube为代表的机器学习方法尝试用深度强化学习或自监督学习来攻克3×3魔方取得了不错但仍有提升空间的成绩最优解率约60%-70%。然而对于更大规模的魔方机器学习似乎一直未能真正涉足。直到最近一项名为CayleyPy的研究提出了一种全新的思路不再让神经网络直接学习“下一步怎么走”的策略而是让它学习估算一个状态距离目标还有“多远”再结合经典的束搜索算法像一位拥有方向感的探险家在迷宫中高效寻路。这个方法的核心创新在于两点一是用“扩散距离”这一图论概念作为神经网络的预测目标二是引入了“多智能体”协作机制。结果令人震惊该方法不仅在3×3魔方上取得了98.4%的最优解率追平了基于模式数据库的专用求解器更是首次成功求解了4×4和5×5魔方且求解路径长度超越了包含上千支队伍的Kaggle Santa 2023挑战赛的最佳结果总和。这标志着机器学习在解决极端大规模组合优化问题上迈出了关键一步。1.1 核心思路化“策略学习”为“距离感知”传统基于学习的魔方求解器如DeepCubeA通常训练一个策略网络Policy Network来直接预测在给定状态下哪个转动动作的“价值”最高。这类似于教AI下围棋需要它评估每一步的优劣。然而在状态空间极其庞大且稀疏的图中学习一个精准的策略非常困难需要海量的探索和精巧的奖励设计。CayleyPy项目则另辟蹊径它回归了更经典的启发式搜索框架。其核心思想是训练一个神经网络让它成为一个高效的“距离估算器”。这个网络不关心具体怎么走只回答一个问题“从当前状态到复原态大概还需要多少步”这个估算值就是启发式搜索中的启发函数Heuristic Function。那么如何定义这个“距离”研究者们采用了“扩散距离”的概念。你可以想象在复原态这个节点上滴一滴墨水让它沿着图的边随机游走。那么一个状态节点首次被这滴“墨水”染色的平均步数就定义了这个节点到复原态的扩散距离。在数学上这等价于从复原态出发的随机游走首次到达该节点的期望步数。这个距离虽然不是精确的图最短路径长度即上帝之数但它与最短路径高度相关并且有一个巨大的优势可以通过简单的随机模拟来生成大量的训练数据。具体来说生成训练数据的过程就是不断地从复原态开始进行随机转动即随机游走并记录下每一步到达的状态以及已经走过的步数。这些状态步数配对就构成了神经网络的训练集。神经网络的任务就是学习从状态向量即魔方的排列编码预测这个步数。一旦这个网络训练好它就能对任何未见过的状态给出一个距离复原态“有多远”的估计值。有了这个距离估计器剩下的路径寻找工作就可以交给成熟的图搜索算法。研究者们选择了束搜索。束搜索是一种宽度受限的启发式搜索算法它不像A*那样需要维护复杂的开放列表和优先级队列而是在每一步扩展时只保留当前看来最有希望的W个节点W称为束宽然后在这些节点的邻居中继续选择最好的W个如此迭代。这种方法在搜索深度和广度之间取得了很好的平衡特别适合与神经网络启发式函数结合因为神经网络的预测可以快速对大量候选节点进行排序和筛选。1.2 方法优势为什么它能处理超大规模问题这套方法之所以能处理像5×5魔方这样状态数超过10^74的庞然大物关键在于其巧妙的设计将计算复杂度分解到了可管理的层面数据生成成本极低训练数据来自随机游走计算开销很小且可以并行生成。这与需要通过与环境交互来收集经验的强化学习方法形成鲜明对比。模型泛化能力神经网络的核心作用是从有限的随机游走样本中学习到整个状态空间上距离函数的“平滑”近似。它不需要见过所有状态就能对未知状态做出合理的距离估计。这解决了状态空间太大无法枚举的根本问题。搜索与学习解耦训练阶段只负责学习距离函数搜索阶段则固定使用训练好的模型。这使得我们可以独立地优化搜索算法如增大束宽W或训练更强的模型而不必像端到端强化学习那样重新进行耗时巨大的策略迭代。多智能体集成由于随机游走的随机性每次训练得到的神经网络对距离的估计会有细微差异。利用这种差异训练多个独立的“智能体”即多个独立训练的神经网络让它们各自进行束搜索求解同一个打乱状态最后选取所有智能体中找到的最短路径。这种集成方法显著提高了找到更短路径的概率相当于让多个拥有不同“方向感”的探险家同时探路取最快到达的那个结果。2. 核心组件深度解析从理论到实现细节要复现或深入理解这项工作我们需要拆解其三个核心组件训练数据生成扩散距离估计、神经网络模型ResMLP以及束搜索路径查找器。每一个环节都有值得深究的设计选择和实操要点。2.1 扩散距离一个巧妙的监督信号扩散距离作为监督信号其有效性建立在两个基础上一是它与最短路径距离强相关二是它易于通过蒙特卡洛方法近似。为什么不用最短路径距离本身最直接的监督信号当然是每个状态到复原态的最短路径长度即精确距离。但对于大规模魔方计算这个精确距离本身就是NP难问题无法为海量训练数据提供标签。扩散距离则提供了一个可计算的替代品。从复原态开始的随机游走到达某个状态所需的平均步数必然会大于或等于最短路径长度因为随机游走会绕路但两者在趋势上是一致的距离复原态越远的状态随机游走到达它的期望步数也越大。实操中的关键参数K_max在生成训练数据时需要设定随机游走的最大步数K_max。这个参数不能太小否则无法生成距离足够远的状态样本也不能太大否则随机游走可能会在状态空间中“迷失”导致生成的状态步数配对中步数k与真实的最短路径距离相关性变弱。论文中根据魔方大小经验性地选择了不同的K_max3×3魔方用26步4×4用45步5×5用65步。这背后是基于对魔方图直径任意两状态间最大最短路径的估计。例如4×4魔方的理论直径猜想为48步那么45步的随机游走足以覆盖绝大部分有意义的距离范围。注意在实现随机游走时论文采用了“非回溯”随机游走即禁止立即回转至上一步的状态。这能提高游走的效率避免在原地附近徘徊从而用更少的步数探索到更远的状态提升了数据质量。训练数据规模与“性能停滞”现象一个反直觉但至关重要的发现是增大训练数据集T对模型性能的提升存在一个饱和点。实验表明当训练样本数达到约80亿8B后继续增加数据量对降低平均求解步数几乎没有帮助。这意味着我们无需追求无限大的数据集在达到某个阈值后计算资源应投入到增加模型容量深度、宽度或增大搜索束宽W上。这为资源分配提供了明确的指导与其耗费大量时间生成更多数据不如用同样的时间训练更多不同的智能体模型利用集成效应来提升性能。2.2 神经网络架构轻量而有效的ResMLP模型采用了一种带有残差块的多层感知机称为ResMLP。其输入是魔方状态的向量化表示。对于n×n×n魔方其状态可以用一个排列向量来表示描述每个色块的位置。经过One-Hot编码后输入维度是6n^26个面每个面n^2个色块。架构细节输入层将One-Hot编码后的稀疏向量通过一个线性层映射到第一个隐藏层维度N1。残差块核心部分由Nr个残差块组成。每个残差块包含两个线性层中间有ReLU激活和批归一化BatchNorm并通过跳跃连接将块输入加到输出上。残差结构有助于缓解深层网络中的梯度消失问题让训练更稳定。输出层最后一个残差块的输出经过一个线性层映射到单个标量值即预测的扩散距离步数k。模型规模与性能的权衡论文实验了不同深度层数和宽度N1, N2的模型。一个关键结论是模型的深度比参数量更重要。即使是一个仅有100万参数的小模型也能在3×3魔方上取得与拥有1470万参数的DeepCubeA模型相近的效果。最终他们选择了一个10层、总参数量约400万4M的模型作为基准。这个模型在性能和计算开销之间取得了良好平衡单个模型训练时间仅需4小时40分钟对比EfficientCube的86小时。训练技巧优化器与损失使用Adam优化器固定学习率0.001损失函数为均方误差MSE直接回归预测的步数k。动态数据集每个训练周期epoch前重新生成100万个新的状态步数样本。这种在线生成数据的方式相当于一种数据增强可以防止模型过拟合到有限的静态数据集并让模型接触到更广泛的状态分布。混合精度训练训练时使用32位浮点数保证稳定性推理搜索时使用16位浮点数可以大幅提升计算速度、减少GPU显存占用这对于需要频繁进行前向传播的束搜索至关重要。2.3 束搜索算法宽度决定成败束搜索是这个方法中最强大也最耗资源的杠杆。其算法流程可以概括为初始化将打乱的初始状态节点放入“束”中此时束内只有1个节点。扩展与评估对于当前束内的所有节点生成它们的所有合法下一步状态邻居节点。利用训练好的神经网络对所有新生成的邻居节点进行距离估计前向传播。排序与筛选将所有评估过的节点包括上一轮保留的节点按照神经网络预测的距离值从小到大排序。剪枝仅保留排名前W个节点作为下一轮迭代的“束”。其余节点被丢弃。终止判断检查当前束中是否包含复原态节点。如果包含则回溯路径输出解序列如果达到最大步数限制如200步仍未找到或束中所有节点都已访问过则判定本次搜索失败。束宽W的决定性影响实验数据清晰地表明平均解长度与束宽W的对数近似成线性负相关。也就是说增大W几乎总能得到更短的解且收益是稳定的。在3×3魔方上将W从2^124096增加到2^24约1678万平均解长从22步多降到了接近最优的20.65步。对于4×4魔方增大W同样显著降低了求解步数。为什么束搜索如此有效在巨大的状态空间中贪婪搜索只选最好的一个节点极易陷入局部最优。而束搜索保留了多个候选路径相当于同时探索多条可能的路线。神经网络提供的距离估计虽然不绝对精确但足以在宏观上区分“更有希望”和“希望不大”的区域。束搜索利用这个粗略的“方向感”进行了一种高效的定向广度探索。W越大探索的广度就越大找到更短路径甚至是最短路径的概率就越高。工程优化哈希去重在扩展节点时会生成大量重复的状态因为从不同路径可能到达同一状态。论文中特别提到他们使用哈希函数来快速识别并去除重复节点。这一步至关重要能极大减少不必要的计算和内存占用。在实现时可以使用类似Zobrist哈希的技术为每个魔方状态计算一个唯一的哈希值并维护一个已访问节点的哈希集合。3. 从零构建求解器实操步骤与参数调优假设我们想针对4×4魔方复现一个基础版本的求解器。以下是基于论文方法的核心步骤和参数选择指南。3.1 环境与数据准备首先需要定义魔方的状态表示和转动操作。对于4×4魔方又称“魔方的复仇”其状态空间巨大必须采用高效的编码。状态表示 一种实用的方法是使用一个长度为966个面 * 4*4的整数数组每个位置存储一个色块编号0-5代表6种颜色。更紧凑的表示是使用排列群但数组表示更直观。转动操作则对应对这个数组进行一系列复杂的、固定模式的元素置换。你需要预先实现所有基本转动如U, U, R, R, F, F, 以及中间层转动等的函数。生成训练数据设定参数K_max 45根据论文num_walks决定总数据量单次游走长度在1到K_max之间随机。从复原状态开始进行非回溯随机游走。记录每一步后的状态数组和当前步数k。将状态数组扁平化并One-Hot编码或直接使用某种特征提取与k配对存入数据集。论文中使用了80亿8B个样本作为饱和点。在实际操作中我们可以从较小的数据量如几百万开始观察模型性能逐步增加。3.2 模型构建与训练使用PyTorch框架构建ResMLP模型。import torch import torch.nn as nn import torch.nn.functional as F class ResidualBlock(nn.Module): def __init__(self, hidden_dim): super().__init__() self.linear1 nn.Linear(hidden_dim, hidden_dim) self.bn1 nn.BatchNorm1d(hidden_dim) self.linear2 nn.Linear(hidden_dim, hidden_dim) self.bn2 nn.BatchNorm1d(hidden_dim) def forward(self, x): identity x out F.relu(self.bn1(self.linear1(x))) out self.bn2(self.linear2(out)) out identity # 残差连接 out F.relu(out) return out class ResMLP(nn.Module): def __init__(self, input_dim, hidden1_dim, hidden2_dim, num_blocks): super().__init__() self.input_layer nn.Sequential( nn.Linear(input_dim, hidden1_dim), nn.BatchNorm1d(hidden1_dim), nn.ReLU() ) self.res_blocks nn.ModuleList( [ResidualBlock(hidden2_dim) for _ in range(num_blocks)] ) # 连接层将hidden1_dim映射到hidden2_dim self.adapter nn.Linear(hidden1_dim, hidden2_dim) self.output_layer nn.Linear(hidden2_dim, 1) def forward(self, x): x self.input_layer(x) x F.relu(self.adapter(x)) for block in self.res_blocks: x block(x) return self.output_layer(x).squeeze() # 输出标量距离估计参数设置参考对应4×4魔方4M参数模型input_dim: 6 * 4 * 4 96 (One-Hot后维度或使用其他特征提取后的维度)hidden1_dim: 1010 (论文Table II, Model No.16)hidden2_dim: 592num_blocks: 4总参数量 ≈ 4M训练循环model ResMLP(...).cuda() optimizer torch.optim.Adam(model.parameters(), lr0.001) criterion nn.MSELoss() for epoch in range(total_epochs): # 每个epoch生成新的100万个样本 states, targets generate_training_data(num_samples1_000_000, K_max45) dataset TensorDataset(states, targets) loader DataLoader(dataset, batch_size1024, shuffleTrue) model.train() for batch_states, batch_targets in loader: batch_states, batch_targets batch_states.cuda(), batch_targets.cuda().float() optimizer.zero_grad() predictions model(batch_states) loss criterion(predictions, batch_targets) loss.backward() optimizer.step() # 可选的验证和保存检查点3.3 束搜索实现束搜索的实现需要仔细处理状态扩展、去重和排序。def beam_search_solve(initial_state, model, beam_width2**20, max_steps200): 使用束搜索求解魔方。 initial_state: 初始状态表示如一个整数数组 model: 训练好的ResMLP模型用于估计距离 beam_width: 束宽W max_steps: 最大搜索步数 model.eval() with torch.no_grad(): # 初始化束元素为状态路径累计成本 # 这里累计成本可以用步数也可以用神经网络估计值。论文中直接用估计值排序。 beam [(initial_state, [], 0)] # (state, path, cost) visited {hash_state(initial_state)} # 哈希集合用于去重 for step in range(max_steps): candidates [] for state, path, _ in beam: # 生成当前状态的所有合法下一步 next_states, moves generate_all_moves(state) for next_state, move in zip(next_states, moves): state_hash hash_state(next_state) if state_hash in visited: continue visited.add(state_hash) # 准备输入 state_tensor preprocess_state(next_state).unsqueeze(0).cuda() # 估计距离 estimated_dist model(state_tensor).item() new_path path [move] candidates.append((next_state, new_path, estimated_dist)) if not candidates: break # 没有新节点可扩展 # 按估计距离排序选择最好的beam_width个 candidates.sort(keylambda x: x[2]) beam candidates[:beam_width] # 检查是否找到解 for state, path, dist in beam: if is_solved(state): return path, dist, step1 # 返回解序列、最终估计距离、所用步数 return None, None, max_steps # 未能在限制步数内找到解关键参数选择束宽W这是最重要的性能旋钮。从2^1665536开始尝试如果GPU内存允许存储大量状态和进行批量前向传播可以逐步增加到2^20约100万甚至更高。W越大解的质量越好但内存和计算开销也越大。最大步数设为200是一个安全值防止搜索无限循环。对于4×4魔方大多数解应在100步以内。多智能体集成要获得论文中的最佳效果需要训练多个例如10-30个独立的模型智能体。每个模型用不同的随机种子初始化并生成不同的训练数据。然后对同一个打乱状态用所有智能体并行进行束搜索最后选择找到的最短路径作为最终解。4. 性能调优与避坑指南在实际操作中你会遇到一系列工程和算法上的挑战。以下是一些关键的注意事项和调优经验。4.1 内存与计算效率的平衡束搜索是内存和计算的双重消耗者。每一步都需要存储W个状态及其路径并评估W * b个候选节点b是每个状态的合法移动数对于魔方是固定的如4×4魔方有24种基本转动。优化策略批量前向传播不要对候选状态一个一个地进行神经网络预测。将一批候选状态例如10万个组成一个张量一次性输入模型进行批量预测可以极大利用GPU的并行计算能力。状态压缩与哈希魔方的原始状态表示如96维数组很占内存。在束搜索中我们不需要完整的色块信息来进行哈希比对和神经网络前传。可以使用更紧凑的表示如将其编码为一个uint64类型的整数通过类似Zobrist哈希的方式。神经网络输入则可以使用一种轻量级的特征提取例如每个面的颜色同质性度量、错位棱块/角块数等而不是完整的One-Hot编码。论文中使用排列向量也是一种紧凑表示。混合精度推理如前所述在搜索时使用torch.cuda.amp进行自动混合精度推理可以节省近一半的显存并加速计算。束的剪枝策略除了保留Top-W还可以加入一些启发式剪枝。例如如果两个候选状态的估计距离非常接近且它们的哈希值显示它们很可能通过对称操作等价可以考虑只保留一个。4.2 模型训练中的陷阱损失震荡与不收敛如果使用动态生成的数据集每个epoch的数据分布都在变化可能导致训练曲线震荡。确保每个epoch的数据量足够大论文用100万并且使用批归一化BatchNorm来稳定内部激活值的分布。如果震荡严重可以尝试略微降低学习率或使用学习率预热Warmup策略。过拟合与泛化尽管数据是动态生成的但如果模型容量过大参数过多而随机游走生成的状态分布不够“均匀”模型可能会过拟合到训练数据中常见的局部模式而对一些罕见状态的距离估计偏差很大。监控在验证集另一批独立生成的随机游走数据上的表现。如果验证损失远高于训练损失考虑减小模型规模或增加数据生成的随机性例如增加K_max。输出值范围神经网络预测的距离值步数k可能随着训练变得很大。可以对目标值进行标准化例如除以K_max让网络预测一个0到1之间的相对距离在搜索时再反标准化。这有助于稳定训练。4.3 束搜索的常见问题与调试搜索停滞束搜索可能陷入“高原”——所有候选节点的估计距离都差不多搜索无法有效推进。这通常是因为神经网络的启发式函数在某个区域失去了分辨力。解决方案一是增加束宽W用广度来弥补启发式函数的不足二是引入轻微的随机性例如以很小概率保留一些非Top-W的节点类似集束搜索中的随机束采样三是检查模型在该区域状态下的预测是否合理可能需要用更多样化的训练数据重新训练模型。找到的解过长即使使用了很大的W找到的解可能仍远长于已知最优解。除了检查W是足够大更要检查神经网络的预测质量。可以抽样检查对于已知最短路径长度的一些测试状态看神经网络的预测值是否与真实长度强相关。如果相关性弱说明模型没有学好需要调整模型结构或训练数据。内存溢出OOM这是增大W时最常见的问题。除了上述的批量处理和状态压缩还可以采用“分层束搜索”或“外部内存搜索”等高级技术。一个简单的折衷是使用时间换空间的策略如果W太大无法一次处理可以将束分成多个批次依次评估和排序虽然慢一些但能突破单次内存限制。4.4 多智能体集成的实战技巧集成是提升性能的利器但如何高效管理多个智能体异步并行求解这是最直接的方式。准备一个任务队列里面是所有需要求解的打乱状态。启动多个进程或线程每个进程加载一个不同的训练好的模型从队列中获取任务进行束搜索求解。最后收集所有结果按解长度排序。智能体选择不是所有智能体都一样好。可以先在一个小的验证集例如100个随机打乱上测试所有训练好的智能体记录它们的平均解长度和求解成功率。优先选择那些平均解更短、成功率更高的智能体参与集成。论文中发现即使某些智能体整体表现不佳但它们可能在少数特定难题上表现出色因此保留一个多样化的智能体池是有益的。资源分配在计算资源有限的情况下可以对“更有希望”的智能体分配更大的束宽W而对表现一般的智能体使用较小的W。这是一种动态资源分配策略。5. 结果分析与未来展望回顾论文中的结果这套方法之所以成功在于它找到了一个巧妙的结合点用廉价的数据训练一个轻量的距离预测模型再用暴力的束搜索去弥补模型预测的不精确性。多智能体集成则进一步利用了随机性的力量将单点失败的风险降到了最低。在3×3魔方上仅用单层MLP模型34万参数和超大束宽W2^24就能达到90.4%的最优解率这已经超越了之前许多复杂得多的深度强化学习方法。而26个智能体的集成则将最优解率推高至98.4%与专门设计的、利用魔方群对称性的模式数据库求解器性能相当。在4×4和5×5魔方上的成功更具里程碑意义。它证明了这种方法具备极强的可扩展性。模型和数据生成过程并不依赖于魔方的具体大小只需要调整输入维度对应色块数和随机游走长度K_max。计算资源的瓶颈主要在于束搜索而束搜索可以通过并行化来加速。个人在实际复现和类似项目中的体会是这套框架的通用性极强。它不局限于魔方任何可以建模为状态转移图特别是凯莱图的规划问题理论上都可以尝试这种方法。例如其他组合拼图如十五数码滑块、某些分子构象搜索问题、甚至是一些抽象代数中的字问题都可能成为其用武之地。关键在于两点一是如何将状态高效地编码为神经网络可以处理的向量二是定义合理的“基本动作”集合来生成状态转移。当然当前方法也有其局限。束搜索在束宽极大时依然非常耗时耗内存。未来的一个方向是探索更高效的搜索算法与神经启发式的结合例如蒙特卡洛树搜索MCTS。另一个方向是设计更强大的神经网络架构例如引入图神经网络GNN来显式地利用状态图的结构信息或许能用更小的模型和更窄的束宽达到相同的效果。最后对于想要亲手尝试的开发者我的建议是从2×2或3×3魔方开始。完整实现4×4魔方的转动生成函数已经是一个不小的工程挑战。先在小问题上验证整个流程数据生成、模型训练、束搜索。确保你的束搜索能找到已知的最优解对于小魔方可以用BFS验证。然后再逐步扩展到更大的问题你会更深刻地体会到每个组件的作用和调参的微妙之处。这个过程本身就像在解一个关于“如何解魔方”的元魔方其乐趣和挑战丝毫不亚于转动那个多彩的立方体。