🤖 machine learning

Efficiently Learning Branching Networks for Multitask Algorithmic Reasoning

本文介绍了分支神经网络，这是一种通过利用凸松弛将任务层次化地划分到树状结构中，从而高效学习多任务算法推理的新型架构，进而显著提升了在各种基准测试中的性能并降低了计算成本。

原作者： Dongyue Li, Zhenshuo Zhang, Minxuan Duan, Edgar Dobriban, Hongyang R. Zhang

发布于 2026-07-15

📖 1 分钟阅读☕ 轻松阅读

原作者： Dongyue Li, Zhenshuo Zhang, Minxuan Duan, Edgar Dobriban, Hongyang R. Zhang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你是一位指挥家，正试图教导一支庞大的管弦乐团，不仅要演奏一首曲子，还要同时演奏三十首不同的复杂交响乐。有些乐曲共享一段旋律；有些则会发生剧烈的冲突。如果你强迫每位乐手在同一张巨大的乐谱上同时演奏所有的曲子，结果就会变成一片嘈 way 的嘈杂。乐手们会感到困惑，音符会混杂在一起，表演也会因此受损。这正是研究人员试图让单个神经网络同时解决许多不同“算法推理”任务（例如在迷宫中寻找最短路径或对数字列表进行排序）时所发生的情况。该论文指出，这种“一刀切”的方法会导致干扰（interference），即一个任务的逻辑（如广度优先搜索）会阻碍另一个任务（如深度优先搜索）的执行，从而导致性能下降。

来自东北大学和宾夕法尼亚大学的研究团队提出了一个聪明的解决方案，称为分支网络（branching networks）。他们不再强迫管弦乐团同时演奏所有曲目，而是构建了一个树状结构的指挥台。

以下是其工作原理：

树状结构： 想象一棵树，树干是表演的起点。随着音乐的推进（逐层展开），树木会分裂成分支。一些分支由相似的任务共享，而另一些分支则为完全不同的任务独立分出。例如，论文发现“广度优先搜索（BFS）”和“贝尔曼-福特算法（Bellman-Ford）”就像是表亲；它们在前几步共享相同的路径，因此可以共享相同的乐手（神经网络层）。但“深度优先搜索（DFS）”是一个叛逆者，它很早就走向了不同的路径，因此拥有自己的分支。
神奇的地图（算法）： 你可能会想：“但如何知道哪些任务属于哪个分支呢？组合方式实在太多了！”作者承认，检查每一种可能性将耗费无穷的时间（复杂度为 $O(k^n L)$ ，这是一个数学上的噩梦）。相反，他们发明了一个快速且聪明的捷径。他们使用一种技术，通过观察“梯度”（你可以将其理解为任务的“音乐指纹”或任务给模型的特定“感觉”），在不进行完整训练的情况下估算两个任务之间的相似度。这使得他们能够以极快的速度绘制出树状图，将复杂度降低到仅为 $O(nL)$。这就像拥有一个智能 GPS，能瞬间识别哪些道路汇合、哪些道路分叉，从而避免了为了检查每条路线而进行的重复行驶。

该论文的实际发现：
研究人员在名为 CLRS 的著名基准测试上测试了这一想法，该测试包含 12 种不同的图算法。他们发现，他们的分支网络（被称为 AutoBRANE）是明显的赢家。

它在准确率上超过了现有的最佳“单网络”尝试 3.7%。
它超过了其他“分支化”尝试 1.2%。
但真正的魔力在于效率：它比之前的最佳方法节省了 48% 的时间（GPU 小时）和 26% 的内存。

他们并未止步于图论。他们还将此方法应用于基于文本的推理任务，使用了大型语言模型（如 Llama 和 Qwen）。即使面对这些庞大的模型（参数量高达 340 亿），他们的方法仍比最强的基线模型在准确率上提升了 3.2%。在涉及 2100 万条边和 500 个不同社区标记任务的大规模测试中，该方法将准确率提升了 28%，并且运行速度比其他分支方法快了 4.5 倍。

该论文排除了哪些可能性：
作者非常明确地指出了哪些做法是行不通的。他们明确反对认为单个扁平化的神经网络可以高效处理所有这些任务的观点。他们证明了，当你试图强迫一个单一网络同时学习所有算法的步骤时，任务之间会产生干扰，导致模型出错。他们还排除了为每个任务都训练一个完全独立的庞大模型的想法，并指出这将需要存储 $n$ 个模型（其中 $n$ 是任务数量），这会导致内存灾难。他们的分支树是“金发姑娘”式的完美方案：既不像单网络那样过于僵化，也不像 $n$ 个独立网络那样臃肿。

他们的结论有多可靠？
论文表现得相当自信，但也措辞谨慎。他们在八种不同的架构和多个数据集上测量了这些结果。他们并非凭空猜测，而是进行了实验。

他们证明了他们的“基于梯度的亲和力（gradient-based affinity）”评分（即他们衡量相似度的方式）可以预测模型的真实性能，误差小于 5%。
他们展示了他们自动学习到的树状结构实际上符合人类对于算法相似性的直觉（例如，将所有基于 DFS 的算法归为一类）。
他们表明这种方法既适用于小型图模型，也适用于巨型语言模型。

该论文表明，这种方法为教导 AI 进行逐步推理开启了一扇新大门，就像人类通过意识到不同谜题共享底层逻辑来学习解决不同类型谜题一样。这并不是一个能瞬间解决一切的魔杖，而是一种高效且具有数学依据的组织多任务混沌的方法。作者甚至提到，尽管他们发现了这些结果，但一个更深层次的问题——即为什么某些算法比其他算法更难学习（例如为什么“普里姆算法/Prim's algorithm”似乎比“BFS”需要更多的训练样本）——仍然是一个有待未来探索的谜团。

技术摘要：高效学习用于多任务算法推理的分支网络

1. 问题陈述

算法推理涉及执行经典算法（如广度优先搜索 BFS、贝尔曼-福特算法 Bellman-Ford）所需的逐步逻辑推理。虽然先前的研究已成功训练神经网络来预测单算法的执行轨迹，但多任务算法推理（同时解决多个不同的算法任务）的挑战仍未得到充分探索。

由于不同算法执行轨迹之间的干扰，为单个神经网络训练多个任务具有内在难度。例如，虽然广度优先搜索（BFS）和贝尔曼-福特算法可能共享中间节点标签，但深度优先搜索（DFS）通常会在特定步骤发生分歧。一个被迫同时预测所有任务的单体网络会遭受冲突梯度的影响，而为每个任务训练单独的网络则会产生极高的内存成本（对于 $n$ 个任务需要 $n$ 个模型）。核心挑战在于设计一种架构，能够在共享相似任务参数的同时，隔离不相似的任务，且不引入穷举所有可能分支结构的计算复杂度。

2. 方法论：AutoBRANE

作者提出了 AutoBRANE（自动分支网络），这是一个构建递归树状分支网络的框架。该架构允许不同的任务子集根据其学习到的相似性，在特定层共享特定的层（模块），并在逻辑发生分歧时分支进入不同的模块。

2.1 核心架构

该网络由 $L$ 层组成。该架构并非采用单一共享模型或 $n$ 个独立模型，而是学习在每一层对 $n$ 个任务进行划分。

分支逻辑： 具有相似中间执行步骤的任务在给定层共享同一个神经网络模块（例如，一个 GNN 层或一个 LoRA 适配器）。
可扩展性： 在 $L$ 层上对最优 $k$ 元树结构进行朴素搜索需要 $O(k^{nL})$ 的复杂度。AutoBRANE 通过采用自顶向下的逐层划分策略，将复杂度降低至 $O(nL)$。

2.2 快速近似程序

关键创新在于一种高效算法，用于在不训练每个可能的子集模型的情况下，确定任务划分。该程序依赖于损失函数的一阶近似：

初始化： 在所有任务上联合训练一个基础模型，以获得初始化权重 $W^{(0)}$ 。
基于梯度的亲和力： 对于任何任务子集 $S_k$ $S_{k}$ ，使用初始化的梯度来近似损失 $\hat{L}(f_{W}(S_k))$ $\hat{L} (f_{W} (S_{k}))$ 。具体而言，该方法将梯度 $\nabla f_{W^{(0)}}$ $\nabla f_{W^{(0)}}$ 视为特征，并通过求解一个凸松弛（逻辑回归）来估计在 $S_k$ $S_{k}$ 上训练模型的性能。
- 层 $l$ 处任务 $i$ 与 $j$ 之间的亲和力得分 $T^{(i)}_{j,l}$ 是通过在包含这两个任务的随机子集上平均预测损失来估计的。
- 这避免了重复训练的需求，将每层的成本降低到 $O(n)$ 。
凸松歇与聚类： 该算法利用这些亲和力得分，通过一个最大化簇内亲和力的凸松弛程序将任务划分为 $k$ 个组。这种划分是从第 1 层到第 $L$ 层递归进行的。

2.2 适用性

该方法是模型无关的。它可以实例化于：

图神经网络 (GNNs)： 通过在每一层为任务簇分配特定的 GNN 模块。
大语言模型 (LLMs)： 通过在预训练 LLM 之上构建低秩适配器（LoRA）的分支结构。

3. 核心贡献

快速近似算法： 作者开发了一种程序，用于识别使用分支网络进行最优参数共享范式的方案。他们通过实验验证，在各种 GNN 和 LLM（参数量高达 340 亿）上，该一阶泰勒展开近似估计出的真实性能误差小于 5%。
自动分支网络设计： 论文引入了一种能够自动学习层次化任务划分结构的架构，提高了多种算法推理任务的平均性能，并可应用于任何基础模型。
效率与性能权衡： 与现有的多任务和分支基准相比，该方法在运行时、内存使用量和准确性之间实现了更优的权衡。

4. 实验结果

4.1 近似准确度

在四种 GNN 架构（MPNN, GAT, Triplet-GMPNN, Edge Transformer）和四种 LLM（Qwen, Llama, CodeLlama）上的实验表明，当权重距离初始化位置在 10% 以内时，一阶泰勒展开近似产生的相对误差对于 GNN <1%，对于 LLM <5%。

4.2 CLRS 基准测试（图算法）

在 CLRS 基准测试的 12 个图算法推理任务上的评估结果如下：

性能： AutoBRANE 的平均准确率比最先进的单多任务网络（Edge Transformer）高出 3.7%。它也比最强的多任务基准（LearningToBranch 和 GradTAG）提高了 1.2%。
效率： 与最佳多任务基准相比，该方法减少了 48% 的 GPU 训练时长和 26% 的内存使用量。
结构发现： 学习到的分支结构揭示了与算法逻辑一致的层次化聚类（例如，将 BFS 和 Bellman-Ford 归为一类，并将基于 DFS 的算法单独聚类）。

4.3 基于文本的推理（LLMs）

在使用了 LoRA 适配器的文本图推理基准测试（CLRS-Text, GraphQA, GraphWiz）上：

AutoBRANE 比最强的多任务基准（GradTAG）在平均测试准确率上提升了 3.2%。
这证明了该方法即使在自然语言设置下，也能捕捉不同层级的任务相似性。

4.4 大规模社区检测

在一个拥有 2100 万条边和 500 个社区标记任务的数据集上：

AutoBRANE 比现有的多任务和分支架构实现了 28% 的准确率增益。
与 LearningToBranch 相比，它将运行时间缩短了 4.5 倍，并且比 GradTAG 使用更少的内存。

5. 意义与主张

论文声称分支架构是多任务算法推理领域的一项关键进展。通过显式建模算法执行轨迹之间变化的相似程度，AutoBRANE 缓解了困扰单体多任务模型的干扰问题。

作者强调，他们的工作强化了理解神经网络中算法可学习性和泛化能力的重要性。他们认为，预测中间算法状态的能力是衡量推理能力的严苛基准，这往往能揭示当前大语言模型（在提示词驱动下，其表现比经过训练的 GNN 差约 65%）存在的差距。

这项工作并不声称解决了所有的推理问题，而是建立了一个可扩展、高效的框架，用于同时学习多个算法任务。它强调，通过神经网络学习算法的“可学习性”与其执行轨迹的复杂度密切相关，而分支结构可以有效地在参数共享与任务特定专业化之间进行权衡。

6. 局限性与未来方向

论文指出，尽管该近似高度准确，但它依赖于网络经过良好初始化且线性近似在权重局部邻域内成立的假设。作者建议的未来工作包括使用这些技术更好地理解算法推理任务的样本复杂度和上下文学习能力，以及探索随机算法的可学习性。