Cappy: Outperforming and Boosting Large Multi-Task LMs with a Small Scorer

This repo contains code of the following paper:

Cappy: Outperforming and Boosting Large Multi-Task LMs with a Small Scorer
Bowen Tan, Yun Zhu, Lijuan Liu, Eric Xing, Zhiting Hu, Jindong Chen
NeurIPS 2023
[arXiv] [Model Card (btan2/cappy-large)]

Getting Started

• Cappy is a pretrained small scorer designed to enhance the performance and efficiency of multi-task LLMs.
• Cappy takes in an instruction and a candidate response as input, and produces a score between 0 and 1, indicating an estimated correctness of the response with respect to the instruction.
• With merely 360 million parameters, Cappy functions either independently on classification tasks or serve as an auxiliary component for LLMs, boosting their performance.
• Also, Cappy enables efficiently integrating downstream supervision without requiring LLM finetuning nor the access to their parameters.
• Furthermore, Cappy is flexible to cooperate with other LLM adaptations, including finetuning and in-context learning, and prompt tuning, offering additional performance enhancement.

Now, Cappy can be loaded with transformers either as a Jax/Flax model or a PyTorch model.

Jax/Flax

from transformers import AutoTokenizer, FlaxAutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained('btan2/cappy-large')
cappy = FlaxAutoModelForSequenceClassification.from_pretrained('btan2/cappy-large')

instruction = """
What label best describes this news article?
Carlyle Looks Toward Commercial Aerospace (Reuters) Reuters - Private investment firm Carlyle Group,\which has a reputation for making well-timed and occasionally\controversial plays in the defense industry, has quietly placed\its bets on another part of the market.
"""
response = 'Business'

inputs = tokenizer([(instruction, response), ], return_tensors='pt')
score = cappy(**inputs).logits[0][0].item()

PyTorch

from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained('btan2/cappy-large')
cappy = AutoModelForSequenceClassification.from_pretrained('btan2/cappy-large')

instruction = """
What label best describes this news article?
Carlyle Looks Toward Commercial Aerospace (Reuters) Reuters - Private investment firm Carlyle Group,\which has a reputation for making well-timed and occasionally\controversial plays in the defense industry, has quietly placed\its bets on another part of the market.
"""
response = 'Business'

inputs = tokenizer([(instruction, response), ], return_tensors='pt')
score = cappy(**inputs).logits[0][0].item()

Below are the scripts to recover the experiments in the paper.

Requirements

Cappy's pretraining and finetuning are both based on Redco, a lightweight tool automating distributed training on both GPUs and TPUs.

To install redco

pip install redco==0.4.13

Sometimes the Jax version needs be adjusted based on your device & environment. Here are some instructions.

To install other requirements,

pip install -r requirements.txt

Pretraining Cappy

Cappy's pretraining uses the code from this example in Redco. We will release Cappy's pretraining data soon.

Evaluting Cappy on PromptSource (zero-shot)

Download Test Data

Following the setting from OPT-IML paper (Section 5.2). We conduct zero-shot evaluation on 11 held-out classification tasks from PromptSource.

bash scripts/download_promptsource_test_data.sh

Running Cappy

python cappy_promptsource.py --model_name_or_path btan2/cappy-large

Results

Baseline results come from OPT-IML paper (Section 5.2).

Boosting FLAN-T5 with Cappy on Big-Bench Tasks

Getting Big-Bench Tasks

We take all 45 generative tasks from Big-Bench in our experiment. The command below process the tasks into .jsonl format.

python scripts/get_bigbench_data.py

The processed datasets can be found in ./bigbench_data, where ./bigbench_data/subset_names.json records all the task names.

Getting FLAN-T5 Outputs

We collect generated outputs (as well as log-likelihoods on evaluation sets) from FLAN-T5 models (from -small to -xxl). They can be downloaded with

bash scripts/download_bigbench_flan_gens.sh

If you want to generate outputs by your self and/or adjust some generation settings, we provide generation code as below that supports distributed inference using multiple GPUs together (in case the model is too large to accomodate on a single GPU, e.g., FLAN-T5-XXL (11B)).

python scripts/bigbench_flan_generate.py \
  --model_name_or_path google/flan-t5-xl \
  --n_model_shards 4

where --n_model_shards refers to the number of shards you want to split the large model into (it's usually the number of GPUs on your device if it's not 1).

Adapting Cappy to boost FLAN-T5

XLA_PYTHON_CLIENT_MEM_FRACTION=.95 python cappy_bigbench.py \
  --model_name_or_path btan2/cappy-large \
  --bigbench_subset_name auto_categorization \
  --bigbench_gen_model flan-t5-xxl \
  --train_size 102400

• XLA_PYTHON_CLIENT_MEM_FRACTION=.95: (In case GPU memory exceeds) adjust the GPU memory pre-allocation to Jax, see here for more details.
• --bigbench_subset_name: the name of subset from Big-Bench (see ./bigbench_data/subset_names.json for all of them).
• --bigbench_gen_model: the FLAN model to be boosted.
• --train_size: the target data size to construct for Cappy's finetuning on the task (collect FLAN outputs, and then truncate or repeat).

See def main(...) in cappy_bigbench.py for all the arguments.

Every sub-task takes 40 mins to run on a single A10G GPU. The result will be logged in ./bigbench_cappy_results/{flan_model}/{subset_name}.json.

Besides, to run all the Big-Bench subsets at once,

python scripts/run_cappy_bigbench.py --cuda_idx 0

Results

To present baseline results, python scripts/present_bigbench_baselines.py

To present Cappy results on all 45 Big-Bench subtasks, python scripts/present_cappy_bigbench_results.py --gen_model_name flan-t5-xxl

The reported numbers on the paper are produced on TPU machines. Here we provide our reproduction results on A10G GPUs in ./bigbench_cappy_results. The gap between them is slight (ΔrougeL <= 0.8).

Acknowledgement

Cappy is Mario's ally throughout Super Mario Odyssey and assists him in various ways. We thank Nintendo for the nice game!