Efficiently Learning Branching Networks for Multitask Algorithmic Reasoning

2025-11-30

AI Summary

To address negative transfer arising from divergent execution paths in multi-task algorithmic reasoning, this paper proposes Branch Neural Networks (BNNs) and AutoBRANE—an automated hierarchical task grouping algorithm. AutoBRANE constructs an affinity score based on gradient similarity and jointly optimizes task grouping via convex relaxation and hierarchical structure search, achieving optimal grouping in *O*(*nL*) time while supporting foundational models including GNNs and LLMs. On the CLRS benchmark, it improves accuracy by 3.7% over the strongest multi-task GNN baseline, reduces runtime by 48%, and cuts memory consumption by 26%; on large-scale graph data, it achieves a 28% accuracy gain and 4.5× inference speedup. The core contribution is the first formulation of algorithm-task-driven dynamic hierarchical collaborative learning, which significantly enhances multi-task compatibility, generalization, and computational efficiency.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Algorithmic reasoning -- the ability to perform step-by-step logical inference -- has become a core benchmark for evaluating reasoning in graph neural networks (GNNs) and large language models (LLMs). Ideally, one would like to design a single model capable of performing well on multiple algorithmic reasoning tasks simultaneously. However, this is challenging when the execution steps of algorithms differ from one another, causing negative interference when they are trained together. We propose branching neural networks, a principled architecture for multitask algorithmic reasoning. Searching for the optimal $k$-ary tree with $L$ layers over $n$ algorithmic tasks is combinatorial, requiring exploration of up to $k^{nL}$ possible structures. We develop AutoBRANE, an efficient algorithm that reduces this search to $O(nL)$ time by solving a convex relaxation at each layer to approximate an optimal task partition. The method clusters tasks using gradient-based affinity scores and can be used on top of any base model, including GNNs and LLMs. We validate AutoBRANE on a broad suite of graph-algorithmic and text-based reasoning benchmarks. We show that gradient features estimate true task performance within 5% error across four GNNs and four LLMs (up to 34B parameters). On the CLRS benchmark, it outperforms the strongest single multitask GNN by 3.7% and the best baseline by 1.2%, while reducing runtime by 48% and memory usage by 26%. The learned branching structures reveal an intuitively reasonable hierarchical clustering of related algorithms. On three text-based graph reasoning benchmarks, AutoBRANE improves over the best non-branching multitask baseline by 3.2%. Finally, on a large graph dataset with 21M edges and 500 tasks, AutoBRANE achieves a 28% accuracy gain over existing multitask and branching architectures, along with a 4.5$ imes$ reduction in runtime.

Problem

• Designing a single model for multiple algorithmic reasoning tasks
• Reducing combinatorial search for optimal branching network structures
• Improving multitask performance while lowering runtime and memory usage

Innovation

• Branching neural networks for multitask algorithmic reasoning
• AutoBRANE algorithm reduces search to O(nL) time
• Clusters tasks using gradient-based affinity scores