Source-linked AI summary

TALM: Tool Augmented Language Models

Aaron Parisi, Yao Zhao, Noah Fiedel

arXiv:2205.12255v1cs.CLcs.AI

TL;DR

Language models need capabilities beyond scale to access changing or private information and state-changing APIs. TALM combines text-to-text tool use with iterative self-play to bootstrap from few demonstrations, and it consistently outperforms non-augmented LMs on knowledge and reasoning tasks. The paper presents this combination as a way for smaller models to outperform larger non-augmented models.

  • Problem

    Scale and retrieval cannot provide access to all live or private data or APIs that read or modify state, while scaling also has substantial practical costs.

  • Method

    TALM combines a text-to-text interface for invoking arbitrary tools with iterative self-play that bootstraps tool-use examples from few labeled demonstrations.

  • Results

    TALM consistently outperforms non-augmented LMs on the knowledge-oriented NQ task and reasoning-oriented MathQA task.

  • Takeaways & Limitations

    Tool augmentation combined with iterative self-play enables smaller models to outperform larger non-augmented LMs, with less dependence on scale.

Abstract

from arXiv · show

Transformer based language models (LMs) demonstrate increasing performance with scale across a wide variety of tasks. Scale alone however cannot enable models to solve tasks that require access to ephemeral, changing, or private data that was unavailable at training time. Many useful tasks may also benefit from LMs being able to access APIs that read or modify state. In this work, we present Tool Augmented Language Models (TALM), combining a text-only approach to augment language models with non-differentiable tools, and an iterative "self-play" technique to bootstrap performance starting from few tool demonstrations. TALM exhibits strong performance on both a knowledge-heavy QA task and a reasoning oriented math task with simple tools. At a given model scale, TALM significantly outperforms non-augmented LMs. We further demonstrate that TALM successfully performs out-of-distribution inferences on both QA and math tasks, where non-augmented LMs fail. Our results suggest that Tool Augmented Language Models are a promising direction to enrich LMs' capabilities, with less dependence on scale.

1 Introduction

Language-model scaling improves benchmark performance but brings high costs and cannot provide live, private, or state-changing information. TALM addresses these gaps by combining tool access through a text-to-text interface with iterative self-play from few demonstrations.

  • Motivation: Larger Transformer language models generally perform better across tasks, but their storage, deployment, training, and inference costs can be prohibitive.Scaling also increases practical burdens for researchers and organizations.
  • Motivation: Memorization-based models can produce incorrect outputs, become outdated as world knowledge changes, and expose undesirable training-data memorization.Retrieval methods reduce dependence on scale but do not address every useful application.
  • Problem: Many applications require access to live or private data, or APIs that read or modify state.Examples include weather and personal calendars.
  • Approach: TALM lets models invoke arbitrary tools with model-generated outputs and attend to tool results when generating task outputs.This differs from approaches where the model only receives environmental queries.
  • Contributions: The paper demonstrates a text-to-text tool API and an iterative self-play technique that bootstraps tool-augmented datasets from few labeled examples.These are presented as the paper’s main contributions.

2 Methods

TALM uses a text-to-text interface to call tools and incorporate their outputs, then iteratively expands tool-use data through self-play. The procedure samples and filters successful tool interactions so later fine-tuning can improve performance.

  • 2.1 Tool Augmented Language Models: TALM generates a tool input, invokes the tool when it emits a delimiter, appends the result, and continues generating the final task output.The interface supports tool use without requiring differentiable tool integration.
  • 2.1 Tool Augmented Language Models: TALM jointly learns to call a tool and generate an answer from the tool result, while remaining compatible with several language-model architectures.The implementation uses the Seq2Seq family because of its fine-tuning performance at modest scale.
  • 2.2 Iterative self-play: Iterative self-play bootstraps tool-use examples from limited demonstrations by exploiting supervised task data and automated output evaluation.The model interacts with the tool API rather than playing adversarially against another model.
  • 2.2 Iterative self-play: Each self-play round fine-tunes on the current tool-use set, samples tool queries and outputs, and adds sequences whose answers match targets within a threshold.The retained sequences become training data for the next round.
  • 2.2 Iterative self-play: Self-play uses random sampling with temperature=1.0 and top-k=40, generates up to N=600 sequences per example, and evaluates with 4-beam decoding.These settings encourage diverse tool invocations during data generation and a single output at evaluation.
  • 2.2 Iterative self-play: The self-play pipeline is a special case of policy-gradient reinforcement learning with the language model as policy and a binary reward signal.The paper notes possible extensions to multihop tool use and broader reinforcement-learning formulations.

3 Results

TALM is evaluated on knowledge-heavy Natural Questions and reasoning-oriented MathQA using retrieval and formula-solving tools. It outperforms non-augmented language models, gains from iterative self-play, and handles changing knowledge and large-number inputs where standard LMs fail.

  • 3 Results: TALM is evaluated on Natural Questions for knowledge-intensive QA and MathQA for general reasoning.NQ uses a retrieval tool over oracle contexts, while MathQA uses a solver tool for formulas.
  • 3.1 Natural Questions: Even the 220M base TALM outperforms the 3B XL LM on Natural Questions.The smaller performance gap among TALM model sizes than between TALM and LM suggests retrieval benefits smaller models more on knowledge-intensive tasks.
  • 3.2 MathQA: TALM significantly outperforms a non-augmented LM on MathQA despite formula validity of approximately 70% in the solver tool.The bootstrap tool-use dataset is a random 10% sample of valid-formula training examples, approximately 2k examples.
  • 3 Results: Performance improves after one self-play round and continues increasing through three rounds across model sizes from 220M to 3B.This compares iterative self-play against training only on the limited bootstrap tool-use examples.
  • 3.4 Out-of-distribution Examples: TALM generalizes to out-of-distribution inputs by handling changing world knowledge and large-number math cases where non-augmented LMs fail.For QA, the BM25 Wiki retriever is replaced with a public search engine; for math, the evaluation targets large-number handling.

4 Conclusion

TALM augments language models with arbitrary tools through a text-to-text interface and iterative self-play, achieving strong performance across knowledge and reasoning tasks. The authors position this combination as a way to improve model capabilities with less dependence on scale and less expensive labeling.

  • TALM combines a text-to-text tool interface with iterative self-play to augment language models using arbitrary tools.The approach targets tasks with few tool-use labeled examples.
  • The interface and self-play technique support exploring additional tools and tasks without requiring expensive data-labeling efforts.
  • TALM consistently outperforms non-augmented language models on both the NQ knowledge task and the MathQA reasoning task.The authors report that ablations identify self-play as important and iterative self-play as providing further gains.
  • The authors conclude that tool augmentation with iterative self-play enables smaller models to outperform larger non-augmented language models.
  • The authors describe tool-augmented language models as a promising direction for enhancing capabilities with less dependency on scale than many contemporary approaches.

5.2 Author Contributions

The author contributions cover project conception, pipeline implementation, experiments, baselines, solver development, research leadership, and paper writing.

  • The section identifies the listed author contributions for the paper.
  • Aaron Parisi designed and implemented the tool-augmentation and self-play pipelines and ran most experiments.He also participated in brainstorming and paper writing.
  • Yao Zhao contributed to brainstorming, experimental setup discussions, and paper writing while implementing NQ and MathQA baselines and MathQA solvers.
  • Noah Fiedel conceived the project, led the research group, and contributed to brainstorming and paper writing.
Loading 2205.12255v1…