Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Token Alignment via Character Matching for Subword Completion

Ben Athiwaratkun, Shiqi Wang, Mingyue Shang, Yuchen Tian, Zijian Wang, Sujan Kumar Gonugondla, Sanjay Krishna Gouda, Rob Kwiatowski, Ramesh Nallapati, Bing Xiang

TL;DR

This paper examines a technique to alleviate the tokenization artifact on text completion in generative models, maintaining performance even in regular non-subword cases.

Abstract

Generative models, widely utilized in various applications, can often struggle with prompts corresponding to partial tokens. This struggle stems from tokenization, where partial tokens fall out of distribution during inference, leading to incorrect or nonsensical outputs. This paper examines a technique to alleviate the tokenization artifact on text completion in generative models, maintaining performance even in regular non-subword cases. The method, termed token alignment, involves backtracking to the last complete tokens and ensuring the model's generation aligns with the prompt. This approach showcases marked improvement across many partial token scenarios, including nuanced cases like space-prefix and partial indentation, with only a minor time increase. The technique and analysis detailed in this paper contribute to the continuous advancement of generative models in handling partial inputs, bearing relevance for applications like code completion and text autocompletion.

Token Alignment via Character Matching for Subword Completion

TL;DR

This paper examines a technique to alleviate the tokenization artifact on text completion in generative models, maintaining performance even in regular non-subword cases.

Abstract

Generative models, widely utilized in various applications, can often struggle with prompts corresponding to partial tokens. This struggle stems from tokenization, where partial tokens fall out of distribution during inference, leading to incorrect or nonsensical outputs. This paper examines a technique to alleviate the tokenization artifact on text completion in generative models, maintaining performance even in regular non-subword cases. The method, termed token alignment, involves backtracking to the last complete tokens and ensuring the model's generation aligns with the prompt. This approach showcases marked improvement across many partial token scenarios, including nuanced cases like space-prefix and partial indentation, with only a minor time increase. The technique and analysis detailed in this paper contribute to the continuous advancement of generative models in handling partial inputs, bearing relevance for applications like code completion and text autocompletion.
Paper Structure (29 sections, 4 figures, 6 tables, 1 algorithm)

This paper contains 29 sections, 4 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Methodology
  4. Partial Token Scenarios
  5. Evaluation
  6. Datasets
  7. Code Generation Tasks
  8. Execution-Based Code Generation Benchmarks
  9. Metrics
  10. Natural Language Tasks
  11. Metrics
  12. Token Alignment Handles Many Partial Token Scenarios
  13. Subword
  14. Code Generation
  15. Natural Language Tasks
  16. ...and 14 more sections

Figures (4)

  • Figure 1: An illustration of token alignment process. We first tokenize the prompt where the last tokens may correspond to partial tokens. We use the backtracked prompt as model input and use the alignment prefix to filter possible tokens at byte or character level. We then mask out non-matched tokens by zeroing their probabilities, which is later used for sampling to select the next token.
  • Figure 2: Effects of token alignment on prompts ending with a subword. The <T> marks the end of the prompt, after which the completion is shown. Without token alignment, the model fails to predict "return" correctly, as the sequence of tokens "re", "turn" is out-of-distribution. Token alignment alleviates this constraint by backtracking to full tokens before "re", then align subsequent generations with the prompt.
  • Figure 3: Partial token scenarios. Note: the underscore denotes a space character.
  • Figure 4: The density estimation of the number of steps during matching process in token alignment.