Paper: Token Alignment via Character Matching for Subword Completion (ICLR-2024)

Reviewer comment (Evaluation):
In terms of prefix-indentation splitting, the results of the proposed method still lagged behind the baseline.

Author response:
This is a good point. We use the baseline datasets where we pre-backtrack so that the prompt does not contain partial token. Therefore, the difference between the baseline and non baseline is as low as possible, in order to control the effect of the prompt. Moreover, we average the results over many curated datasets to understand the aggregate behavior.

Reviewer comment (Evaluation):
The method introduced a certain time delay.

Author response:
Good point. One way to use in-context learning is to make the model repeat the entire input and generate full tokens that match the partial tokens instead. However, if the input is long, this will amount to high extra latency. Therefore, we do not explicitly consider this approach. 

It is possible that we do in-context learning in such a way that makes the model output only a snippet shortly before the end of the prompt. Such method is likely quite complicated and requires processing to extract out the right portion. 

Our method presents a clean way to align tokens natively by guiding the token selection with the probabilities that the model already predicts, but masking out tokens that do not align. We believe that this is the cleanest and simplest approach.

Reviewer comment (Experiments):
How did the proposed method perform in multilingual or cross-lingual scenarios?

Author response:
the effectiveness of token alignment itself is independent of the cross or multi-lingual setting, holding other factors constant. If cross lingual entails rare language, it would require models to have seen sufficient amount of such language to be able to generalize, even with full token scenario. Given such cases where models can output meaningful probability distribution over tokens in each step, our method helps when there is a presence of partial token.

Rare language may entail higher backtracking B. In such case, a tokenizer schema where we know which token corresponds to the beginning of pre-token can be helpful.

Reviewer comment (Evaluation):
In which areas did "token alignment" excel? Were there scenarios or applications where other methods might be more appropriate?

Author response:
This is a great question. The token alignment is very suitable for 'completion' mode. For instance, if we use language models to help write word document, or write code within the code editor. This is in contrast to the Question - Answer mode where the model does not have to output text that is a continuation of the prompt, but would perhaps start in a different line and start a new sentence / piece of code instead.

Reviewer comment (Experiments):
For long texts or texts exceeding the model's maximum input length, did the "token alignment" method maintain its effectiveness? In such cases, was there a need to adjust or optimize the method?

Author response:
Our method works for different input length. In short, if using the model without token alignment can support L input tokens, using token alignment also supports ~ L input tokens as well (give or take the number of backtrack tokens B where we use around 1-3, as demonstrated in the paper)

Reviewer comment (Writing):
The paper is somewhat unclear about the method of backtracking multiple tokens (B in Algorithm 1), so I would ask the authors to clarify this. The introduction mentions “Our approach involves backtracking to the last complete tokens”, which is how I understood the method initially. But later it is suggested that the method backtracks multiple tokens (B = 3 tokens), and not just one token back.

Author response:
This is a very good question that we will emphasize further to make it clear! For different tokenizers, it can be the case that there can be no 'marker' if that token is the start of a pre token. In this case, B needs to be pre-set to be a certain number. However, in some tokenizer, such as SentencePiece, we can have such marker for pre-token, which allows us to backtrack with different B dynamically. (B can be 1 for some input and 2 or other inputs)

Backtracking 3 tokens can be better in certain case. For example, if we have a very long token that gets broken up into partial tokens. For instance, if the full token is "abcdefghijk" (synthetic example here) but the prompt is "abcdefgh" -- instead of being a full token, this prompt likely gets broken up into something like 'abcd', 'ef', 'gh'. If we do not roll back for 3 steps in this case, then the prompt that the model sees still corresponds to partial tokens that got broken up. This would still constrain the model due to the artifact of tokenization.

However, if we correctly rolls back to before "abcdefgh", then the model can see that the token "abcdefghijk" matches in prefix with the text "abcdefgh" -- so our token alignment method will do the job correctly. 

That being said, B=3 typically is not required for most prompts, but it does not hurt to roll back in order to a full token that may be broken into up to 3 tokens as well. However, as the reviewer nicely pointed out, it gives the model less context. Therefore, it is a balance between giving the model context and precise alignment. The alignment process in itself may be seen as guiding the model in a way, however.

Reviewer comment (Evaluation):
I am not sure why backtracking 3 tokens would work better than 1, as found in the ablation study. Wouldn’t that give the model less context unnecessarily, given that 1 token back is all we need to get to a complete token?

Author response:
This is a very good question that we will emphasize further to make it clear! For different tokenizers, it can be the case that there can be no 'marker' if that token is the start of a pre token. In this case, B needs to be pre-set to be a certain number. However, in some tokenizer, such as SentencePiece, we can have such marker for pre-token, which allows us to backtrack with different B dynamically. (B can be 1 for some input and 2 or other inputs)

Backtracking 3 tokens can be better in certain case. For example, if we have a very long token that gets broken up into partial tokens. For instance, if the full token is "abcdefghijk" (synthetic example here) but the prompt is "abcdefgh" -- instead of being a full token, this prompt likely gets broken up into something like 'abcd', 'ef', 'gh'. If we do not roll back for 3 steps in this case, then the prompt that the model sees still corresponds to partial tokens that got broken up. This would still constrain the model due to the artifact of tokenization.

However, if we correctly rolls back to before "abcdefgh", then the model can see that the token "abcdefghijk" matches in prefix with the text "abcdefgh" -- so our token alignment method will do the job correctly. 

That being said, B=3 typically is not required for most prompts, but it does not hurt to roll back in order to a full token that may be broken into up to 3 tokens as well. However, as the reviewer nicely pointed out, it gives the model less context. Therefore, it is a balance between giving the model context and precise alignment. The alignment process in itself may be seen as guiding the model in a way, however.

Reviewer comment (Experiments):
The problem of partial tokens could be presented more accurately in some instances. For example, at the end of page 3 the authors suggest that subword tokens of linguistic words like “banana” (“banan” or “bana”) could lead to issues. I think it would be best to cite work on this, or prove this experimentally? Linguistically unsound subword tokenisations are present in every LLM, yet they work well in most cases, presumably because there is enough data for models to learn how these subwords combine to form words (e.g. on https://platform.openai.com/tokenizer the word “bananas” is segmented “ban-anas”.)

Author response:
Thank you for pointing this out. We will add additional examples to make it clear. For the current draft, examples given in Figure 1, 2, 3 and in Section 4 illustrates various scenarios. For example, in figure 3, "_ _ _" + "_ return" are the fully tokens corresponding to 4 spaces + "return". In the scenario where the model is presented to " _ _ _ _" as the end of a prompt, then it constitutes a partial token issue. 

It is true that unsound subword tokenization are prevalent. However, once the prompt breaks from the patterns that the models observe during training, then problem arises. Such discrepancies happens when the end of the prompt is cut off before a full token. 
For instance, 'banana' could be composed of two tokens 'ban' and 'ana'. If the prompt ends with "monkeys eating yellow ban", the the most likely next token is 'ana' and there is nothing preventing the model to predict such a token. 

However, if the prompt ends with "monkeys eating yellow bana" and the models have seen 'ban' + 'ana' for every semantic instance of 'banana' during training, then the model will most likely not predict 'na' after this prompt.

Reviewer comment (Novelty):
How does your work differ from this paper? https://aclanthology.org/D19-1507.pdf It seems like they are using the same backtracking algorithm (see Section 4), but testing it with different models and datasets.

Author response:
Thank you for pointing out this work.  The paper mentioned uses beam search together with approximate marginalization. For modern language models, beam search is not typically used but popular methods are top-p (nucleus sampling) and top-k. Our proposed approach is designed to be compatible with top-p and top-k and works seamlessly with such sampling approaches. We will add this to related work and contrast the difference, however. Thank you again!

Reviewer comment (Reproducibility):
Will you release your evaluation datasets publicly?

Author response:
Yes.

Reviewer comment (Reproducibility):
What model’s tokeniser was used for the examples in Figures 1-3? If these examples were constructed by hand, this should be mentioned.

Author response:
These are constructed by hands. We will point it out.

Reviewer comment (Novelty):
It is not clearly described the background knowledge needed to motivate and position the proposed method in the literature.

Author response:
We describe the motivation in the introduction as well as accompanying figures where if prompt gets tokenized where the last token corresponds to a partial token, this corresponds to mismatch between training time versus inference time. The scores where we use token alignment versus without using demonstrates that partial token can hurt language model's generation abilities significantly.

Reviewer comment (Writing):
It is not clearly described the alignment task and the relation to code and NLP.

Author response:
The token alignment task is applicable to *any* text completion task -- we illustrate examples in Figure 1 and 3 and provide experiments for both text and code throughout the paper.

Reviewer comment (Evaluation):
A possible extra contribution can be the addition of a statistical significant test of the results.

Author response:
We make note on the statistical significant -- due to the scale of the evaluation we do not have them by default in this draft.

Reviewer comment (Writing):
Please elaborate in the background used to develop the proposed method.

Author response:
We describe the motivation in the introduction as well as accompanying figures where if prompt gets tokenized where the last token corresponds to a partial token, this corresponds to mismatch between training time versus inference time. The scores where we use token alignment versus without using demonstrates that partial token can hurt language model's generation abilities significantly.

Reviewer comment (Experiments):
Is the alignment task based on fine-tuning a model or in-context-learning? Could you elaborate more on differences and benefits, compared to your method.

Author response:
The token alignment is applicable for both pretrained and finetuned model. In our case, we use pretrained model. We also experimented with few-shot prompting which is the 'in-context learning' scenario (Table 3).

Reviewer comment (Experiments):
Could you elaborate on the selection and importance of hyper-parameters? such as backtracking B.

Author response:
We experimented with different values of B (number of backtrack tokens) -- Table 6 in appendix illustrates our extensive result. For the main paper, we use B=3. For B=1 or 2, we observe small differences.