Technical Summary: Learning to Extrapolate to New Tasks: A Relational Approach to Task Extrapolation

1. Problem Statement

Modern learning systems, including large-scale foundation models, excel at interpolation within the support of their training distribution but struggle to generalize to tasks lying outside this support (out-of-support or OOS). This failure persists even in structured settings, such as predicting functions with parameters beyond the training range, increasing compositional depth, or recombining known primitives in unseen ways. Current models often rely on heuristics tied to the training distribution rather than internalizing the generative mechanisms of a domain, leading to saturation at the boundaries of their experience.

The paper defines task extrapolation as generalization to target tasks where the task-defining structure itself lies outside the training support. Unlike standard out-of-distribution or length generalization where test inputs are unions or extensions of training inputs, task extrapolation requires reasoning about the transformation of the task structure itself. The authors posit that while extrapolating to arbitrary unseen tasks is ill-posed, many real-world tasks are systematically related through shared, unknown structures.

2. Methodology: The Relational Task Extrapolator (RTE)

The core proposal is the Relational Task Extrapolator (RTE), an algorithm that reframes extrapolation as a transductive problem. The central observation is that extrapolation is inherently relational: success depends on learning how tasks transform into one another.

2.1. Core Mechanism

RTE decomposes an unseen target task $f_{\theta^*}$ into two components observed during training:

1. An Anchor Task ($f_{anc}$): A familiar task within the training distribution.
2. A Transformation ($\phi$): A relative descriptor linking the anchor to the target (e.g., a parameter shift, a recursive extension step, or a composition operator).

Instead of learning a direct mapping from task parameters to outputs, RTE learns a relational operator $\Psi$ that maps an anchor-transformation pair to predictions for the target task:
$$\hat{y} = \Psi(x, f_{anc}, \phi)$$
This reduces the difficult out-of-support (OOS) extrapolation problem into a more tractable out-of-combination (OOC) problem. Both the anchor and the transformation are individually observed in the training data, even if their specific combination is novel.

2.2. Three Extrapolation Regimes

The authors instantiate RTE across three canonical regimes:

• Parameter Extrapolation (Continuous): Tasks are indexed by continuous parameters. Transformations are shifts in the parameter space ($\phi = \Delta\theta$). The operator learns how function curves deform under parameter shifts.
• Length Extrapolation (Recursive): Tasks are organized by hierarchical complexity (e.g., polynomial degree). The transformation is an inductive step extending a task of complexity $L-1$ to $L$ (e.g., adding a new term).
• Compositional Extrapolation (Combinatorial): Tasks are formed by combining primitives (e.g., $f \circ g$). One component serves as the anchor, and the other as the transformation argument.

2.3. Inference and Search

A practical challenge is selecting the anchor and transformation when task descriptors are latent. RTE addresses this via a proxy geometry (using Task2Vec embeddings derived from the Fisher Information Matrix) and a two-phase process:

1. Training: The model learns $\Psi$ by minimizing prediction error on pairs of training tasks $(f_i, f_j)$ where the transformation $\phi_{ij}$ is known or inferred.
2. Inference:
   • Continuous Regime: Uses a geometric shortcut. The target context is embedded, and the nearest neighbor in the training library is selected as the anchor. The transformation is inferred as the vector difference in the proxy space.
   • Discrete Regimes (Length/Composition): Employs an amortized search. A decomposer network proposes candidate anchor-transformation pairs, which are then re-ranked by evaluating their likelihood or reconstruction error on the sparse target context.

2.4. Extension to Sequence Models (LLMs)

RTE is extended to Large Language Models (LLMs) by framing task extrapolation as a supervised fine-tuning (SFT) problem. The relational operator is instantiated as a fine-tuned instruction-following model. At test time, the model is prompted with an anchor task (demonstrations), a transformation description (e.g., "shift bit index 3"), and a query. The model learns to apply the transformation logic to the anchor to solve the target. Inference involves a discrete search over candidate decompositions, scoring them based on the model's likelihood on the target context.

3. Key Contributions

• Relational Task Extrapolation (RTE): An algorithmic framework that leverages inter-task relationships to reduce OOS extrapolation to OOC problems.
• Extrapolation Regimes: Formal instantiation of RTE across parameter, length, and compositional extrapolation settings.
• Training and Inference Protocols: Development of protocols using Task2Vec embeddings for proxy geometry and formulating extrapolation as an optimization problem. For LLMs, the use of likelihood on sparse context as a proxy reward for search and verification without extrinsic signals.
• Empirical Validation: Comprehensive evaluation on synthetic function prediction tasks and LLM benchmarks (Sparse Parity, CodeIO).

4. Experimental Results

The authors evaluate RTE against standard inductive baselines and oracle variants (which have access to ground-truth decompositions).

• Synthetic Function Prediction:

  • Parameter Extrapolation: Inductive models fail catastrophically (MSE $\approx 10^5$) on out-of-range parameters due to boundary saturation. RTE reduces error significantly (MSE $\approx 10^2$), bridging the gap toward oracle performance.
  • Length Extrapolation: RTE substantially outperforms naive baselines in predicting high-degree polynomials (MSE 0.37 vs 0.58), though a gap remains compared to the oracle due to imperfect anchor retrieval.
  • Compositional Extrapolation: RTE achieves lower aggregate MSE (0.29) than baselines (0.39) on unseen function compositions, successfully inferring underlying primitives.

• LLM Experiments:

  • Sparse Parity (Length): RTE improves accuracy from 52.86% (Baseline) to 66.07% on subset sizes unseen during training, recovering recursive structure.
  • CodeIO (Composition): RTE nearly doubles accuracy (45.3%) compared to Few-Shot (19.8%) and Chain-of-Thought baselines (30.2%), effectively constraining the model to valid algorithmic steps.

5. Significance and Limitations

The paper claims that RTE provides a principled approach to generalization by exploiting the relational structure of task spaces. It demonstrates that models can extrapolate beyond their training support if they learn to decompose new tasks into known components and transformations.

Limitations and Future Directions:

• Inference Cost: RTE requires test-time computation to identify optimal decompositions, unlike single-pass inductive models.
• Manifold Density (Reachability): RTE assumes the target task is reachable via a transformation existing in the training distribution. Completely alien mechanisms remain unreachable.
• Identifiability: In discrete regimes, the current approach relies on ground-truth relational metadata during training. While the authors propose self-labeling via Expectation-Maximization to relax this, inferring latent relational structure from scratch remains an open challenge.
• Multi-Step Extrapolation: While initial results show success with iterative shifts, chaining multiple transformations without compounding errors requires further investigation.

The authors emphasize that RTE is not a universal black-box solution but is specifically designed for domains possessing an underlying relational structure, offering a pathway for models to internalize and reuse generative mechanisms rather than relying on surface-level heuristics.