Task-Adaptive Tokenization: Enhancing Long-Form Text Generation Efficacy in Mental Health and Beyond (2023)

Abstract: We propose task-adaptive tokenization as a way to adapt the generation pipeline to the specifics of a downstream task and enhance long-form generation in mental health. Inspired by insights from cognitive science, our task-adaptive tokenizer samples variable segmentations from multiple outcomes, with sampling probabilities optimized based on task-specific data. We introduce a strategy for building a specialized vocabulary and introduce a vocabulary merging protocol that allows for the integration of task-specific tokens into the pre-trained model’s tokenization step. Through extensive experiments on psychological question-answering tasks in both Chinese and English, we find that our task-adaptive tokenization approach brings a significant improvement in generation performance while using up to 60% fewer tokens. Preliminary experiments point to promising results when using our tokenization approach with very large language models.

## introduction during a time when mental health support is quickly growing @xcite ), text generation techniques have been identified as potentially useful tools to assist mental health professionals (mhp) and provide mental health support to those who share their struggles online @xcite @xcite . with the help of such generation tools, counselors and social workers alike can increase their work efficiency and offer timely and expert feedback to those in need, especially for those in under-resourced areas where access to mental health services is challenging @xcite @xcite . the task of psychological question-answering (psyqa) is to generate a supportive response to a 1 our work is available at github.com/michigannlp/t ask-adaptive_tokenization. text token seq input ids sampling unique decoder ? isolation ? ... having ... social a sense of having not traditional tokenizer task-agnostic vocabulary subword segmentation 1. under-represented tokens 2. long sequence task-adaptive tokenizer task-specific vocabulary variable segmentation 1. task-specific tokens 2. reduced sequence length ... a sense of purpose 亲密 a sense of e.g., "social isolation and not having a sense of purpose in life have been linked to mood disorders" unique seq (16 tokens) [ ha ] [ ving ] [ a ] [ sense ] [ of ] [ pu ] [ rp ] [ social ] [ iso ] [ la ] [ tion ] [ and ] [ not ] seq 1 (9 tokens) [ social isolation ] [ and ] [ not having ] seq 2 (5 tokens) [ social ] [ isolation ] [ and ] [ not ] [ ha ] [ ving ] [ a sense of ] [ purpose ] [ in life ] [ a sense of purpose ] [ in life ] seq n ... social isolation ... social ... ... [ ose ] [ in ] [ life ] help-seeking post, with an important requirement that the language used in the response aligns with the style used by mhps @xcite . the task poses significant challenges due to its domain-specific terminology, text structure, and language style, which are often underrepresented in the pre-training data. recent research has demonstrated that the linguistic patterns of mhps align with established counseling principles and behaviors @xcite , thus exacerbating the difficulty of the task as pre-training data is often sourced from laypeople @xcite . despite recent progress in large-scale language models (llms) @xcite @xcite , we found that few-shot abilities of these models lag behind in mental health applications due to this misalignment with the mhp linguistic behaviors and style (see appx a). this requires us to tune models to the downstream domain. however, much of the current work focused on fine-tuning @xcite still yields unsatisfying results due to the scarcity of professional data, which is often insufficient to capture all the task-specific intricacies. in this paper, we introduce task-adaptive tokenization as a strategy to adapt a model to a downstream task. rather than focusing the adaptation effort on post-training fine-tuning, which has been the typical approach used in recent work @xcite @xcite , we propose to make task-specific adjustments in the way the language is segmented during the tokenization stage (see our motivation elaborated in sec 3). thus, we propose task-adaptive tokenization, which is built specifically from the task data and equipped with variable text segmentation, yet can be seamlessly integrated into off-the-shelf language models. to illustrate the potential of task-adaptive tokenization, we show an example in fig 1 . as seen in this example, our tokenization process allows for the inclusion of domainspecific terminology such as [_social_isolation] and [_a_sense_of_purpose] in the vocabulary. a model trained with task-specific tokenization is now able to generate these tokens through learned preference, which we show can lead to significant performance improvements. we confront three major challenges when designing tailored tokenization strategies. first, the creation of a task-specific tokenization vocabulary must be performed through an automatic process due to the labor-intensive and time-consuming nature of manual selection. integrating this taskspecific vocabulary seamlessly with pre-trained models poses is challenging, and it requires techniques to fuse the task-specific vocabulary with the pre-trained vocabulary and fine-tune the resulting model accordingly. lastly, we need to address the poor representation of newly added tokens in task-specific vocabularies that were not learned during the pre-training phase. in this paper, we propose task-adaptive tokenization as a method for enhancing text generation in specialized tasks such as psyqa. this paper makes three main contributions. (1) building on insights from cognitive linguistics @xcite , we advocate for using taskspecific data and the developing variable segmentation for a downstream vocabulary as a pre-step for creating a task-adaptive tokenizer. (2) we construct a protocol for merging task-specific and pretrained vocabularies, allowing for fine-tuning inputs to be sampled from multiple tokenization results. (3) we propose a simple yet effective initialization mechanism to alleviate the difficulty of learning representations for new tokens unseen during pre-training. through thorough experiments on the psyqa task in both chinese and english, we demonstrate the significant improvements achieved by our task-adaptive tokenization approach. notably, we achieve these enhancements while utilizing 60% fewer tokens compared to expressing equivalent content length. in addition, we show that our tokenization brings significant improvement in 7b llama models, which suggests that our method is effective regardless of the model size and can unlock additional performance despite the booming era of llms. ## the psyqa task the goal of the psyqa task is to generate a supportive response to the help-seeker via responding to their post, where an essential requirement is to imitate the use of language that is characterized as professional by previous work @xcite . figure 3 shows an example of a question-answer pair in this dataset. posts and responses are often extensive, with help-seekers providing detailed accounts of their experiences, and those offering assistance providing comprehensive views, including emotional comfort, indepth analysis, and various suggestions. the formal definition of the task is as follows: given a question q, a description d, and keywords k, let context c denote an integral description of q, d, and k; @xmath0, c 2 , ..., c m ) is the sequence by segmenting the context c. we aim to generate a response r (a sequence of @xmath1, r 2 , ..., r n )) corresponding to the context c. ## motivation our motivation builds on arguments stemming from cognitive science, where (1) a clear distinction is being made between the vocabulary used to interpret language versus the vocabulary used for language production; and (2) there is evidence for increased efficiency in speech behavior stemming from individual segmentation granularities. these arguments are further expanded by optimization and efficiency goals, which are better achieved in the presence of flexible segmentation. ## task-adaptive tokenization our approach to task-adaptive tokenization consists of three main steps: ## experiments we assess the effectiveness and efficiency of our task-adaptive tokenization on the psyqa task through several automatic and manual evaluations. ## conclusion in this paper, we proposed task-specific tokenization as a way to adapt the generation pipeline to the specifics of a downstream task. we demonstrated the efficiency and improved long-form quality of generation for the domain of mental health, where we specifically addressed the task of psychological question answering. our tokenizer leverages the specifics of the downstream task data, while still retaining the ability to integrate into existing pre-trained models. through extensive experiments, we demonstrated the ability of taskadaptive tokenization to enhance both the effectiveness and efficiency of long-form generation. we believe our work is particularly useful in the era of large language models (llms), as the proposed task-adaptive tokenizer can lead to significant improvements while being domain and model agnostic. based on our findings, we suggest plugand-play tokenization for llms when performing specific generation tasks.