Task Agnostic Architecture for Algorithm Induction via Implicit Composition

TL;DR

The work tackles building a single task-agnostic architecture capable of learning and composing arbitrary algorithms across domains. It proposes a Transformer-like processor that operates on a state $S$ containing executable skill blocks, memory, and registers, and executes via a CPU-like loop to iteratively update $S$. Algorithm composition is formalized as a recursive mixture: $F_C^{St}(S^{i}, t) = \begin{cases} S^{i+1}, & t = 0 \\ F_C^{St}(S^{i+1}, t-1), & \text{otherwise} \end{cases}$, with $S^{i+1}=\sum_{s\in S} F_{comp}(S^{i}, s, t, W^{\alpha})$ and $F_{comp}(S^{i},s,t,W^{\alpha})= \alpha F^{St}(s)(W^{in}(s, S^{i},t)S^{i})$. The mixture coefficients $\alpha$ are produced by a language-model-based component $F_{\alpha}$, enabling guided selection of sub-skills, and training can proceed via gradient-based optimization or discrete search. By adopting a hierarchical curriculum of skills, the framework aims to enable scalable, interpretable, and reusable sub-solutions that can robustly generalize while mitigating forgetting.

Abstract

Different fields in applied machine learning such as computer vision, speech or natural language processing have been building domain-specialised solutions. Currently, we are witnessing an opposing trend towards developing more generalist architectures, driven by Large Language Models and multi-modal foundational models. These architectures are designed to tackle a variety of tasks, including those previously unseen and using inputs across multiple modalities. Taking this trend of generalization to the extreme suggests the possibility of a single deep network architecture capable of solving all tasks. This position paper aims to explore developing such a unified architecture and proposes a theoretical framework of how it could be constructed. Our proposal is based on the following assumptions. Firstly, tasks are solved by following a sequence of instructions, typically implemented in code for conventional computing hardware, which inherently operates sequentially. Second, recent Generative AI, especially Transformer-based models, demonstrate potential as an architecture capable of constructing algorithms for a wide range of domains. For example, GPT-4 shows exceptional capability at in-context learning of novel tasks which is hard to explain in any other way than the ability to compose novel solutions from fragments on previously learnt algorithms. Third, the observation that the main missing component in developing a truly generalised network is an efficient approach for self-consistent input of previously learnt sub-steps of an algorithm and their (implicit) composition during the network's internal forward pass. Our exploration delves into current capabilities and limitations of Transformer-based and other methods in efficient and correct algorithm composition and proposes a Transformer-like architecture as well as a discrete learning framework to overcome these limitations.

Figures (2)

• Figure 1: Diagram in part (a) provides a simplified representation of our proposed architecture. The input is a state made up of blocks that store assembly code, memory for data and registers and stack for execution. The code stored in the state is then executed iteratively via a Transformer 'CPU' and the results are written back into the state which is outputted. Part (b) provides a minimal example of the core components that make up the state.
• Figure 2: Part (a) describes the execution of a single assembly command, $S^{i+1} = E(S^{i})$, S is the state and E is the Transformer execution block (more details in Figure \ref{['fig:input_struct']}) Part (b) describes the hierarchy of skills to compute the linear layer, $Z = w^{T} + b$.