pastalib 0.1.2

PyTorch Implementation for PASTA, A Post-hoc Attention Steering Approach that enables users to emphasize specific contexts for LLMs.

🍝 PASTA: Post-hoc Attention Steering for LLMs 🍝

Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs (Zhang et al. 2023).

PASTA allows a user to improve LLM controllability by simply emphasizing part of the prompt (e.g. the instruction) that the LLM should focus on. It requires no changes to LLM weights and no increase in inference time.

Quickstart -- use PASTA for improved inference

1. Install pastalib:

pip install -e .
# Alternatively,  
# pip install pastalib
# pip install git+https://github.com/QingruZhang/PASTA

2. Initialize PASTA along with a pre-trained LLM.

from pastalib.pasta import PASTA 
from transformers import AutoModelForCausalLM,AutoTokenizer

name = "huggyllama/llama-7b"
model = AutoModelForCausalLM.from_pretrained(name)
tokenizer = AutoTokenizer.from_pretrained(name)

# Select the attention heads to be steered, 
# following the format of {'layer_id': [head_ids]}: 
head_config = {
    '16': [16, 21],
    '12': [23, 26, 0],
    '10': [23, 7, 22],
}

# Initialize the PASTA steerer
pasta = PASTA(
    model=model,
    tokenizer=tokenizer,
    head_config=head_config, 
    alpha=0.01, # scaling coefficient
    scale_position="exclude", # downweighting unselected tokens
)

3. Select specific input spans to emphasize, and then run inference as normal.

# Model Input 
texts = ["Mary is a doctor. She obtains her bachelor degree from ... Return her occupation in JSON format."]

# ===== Before =====
# inputs = tokenizer(texts, return_tensors="pt")
# outputs = model.generate(**inputs)
# ---------------------
# ["Mary is a doctor"]  # returns answer in the wrong format

# ===== After =====
inputs, offset_mapping = pasta.inputs_from_batch(texts)
# User highlighs specific input spans
emphasized_texts = ["Return her occupation in JSON format"]
# PASTA registers the pre_forward_hook to edit attention
with pasta.apply_steering(
    model=model, 
    strings=texts, 
    substrings=emphasized_texts, 
    model_input=inputs, 
    offsets_mapping=offset_mapping
) as steered_model: 
    outputs = steered_model.generate(**inputs)
# -------------------------------
# ["{'occupation': 'doctor'}"]  # returns answer in the correct format

Additional Note

1. pastalib works with any models that apply causal attention by summing up query-key inner product with attention mask and can be applied to LLMs in a plug-and-play manner. For example, LlamaForCausalLM and GPTJForCausalLM from transformers, whose attention moduels apply attention masks following torch.matmul(query, key) + attention_mask. However, pastalib currently only supports LLAMA, LLAMA-2 and GPT-J (more models in progress!).

2. We provide different options of head_config for LLAMA-7B and GPT-J in the folder of config/head_config, including multi-task, task-agnostic and task-specific settings. Please see detailed discussion in our paper.

Overview

The overview of this repo is as follows:

• pastalib: contains the source code of PASTA libary, which can be applied to models from huggingface transformers.
• evaluation: consists of evaluation pipelines for different tasks, including data/model preprocessing and task evaluators/metrices.
• config: includes the head_config for steering attention modules of LLAMA-7B and GPT-J with PASTA.
• scripts: consists of running scripts of four tasks: JSON Formatting, Pronouns Changing, Bias in Bios, and CounterFact.

Evaluation

The evaluation pipeline are mainly refactored from REMEDI repo. Please see more details there.

Environment Setup

Set up the environment with the following commands:

conda create -n pasta python=3.10 
pip install -r requirements.txt 
pip install -e . 

python -m spacy download en_core_web_sm
python -W ignore -m nltk.downloader punkt cmudict

By default, the preprocessed datasets, models, and results are saved in the local directory of ./data, ./models, and ./results. You can change the directory of their by setting the environment variables:

export CM_DATA_DIR=<data path> 
export CM_MODELS_DIR=<models path>
export CM_RESULTS_DIR=<results path>

Dataset Setup

1. For CounterFact, our scripts can automatically download the dataset.

2. For Bias in Bios, we cannot release the dataset without the authorization. The dataset must be downloaded with the offical release. After downloading the data examples into the BIOS.pkl file, you can run the following scripts:

python reformat_dataset.py \
--biasbios_raw_path <path of BIOS.pkl> \
--biasbios_save_file biasbios.json 

Then, the preprocessed biasbios dataset file will be saved in CM_DATA_DIR/biasbios.json.

Evaluation

Choose any head_config files from config/head_config and evaluation the performance of PASTA with the following command.

JSON Formatting

python -m scripts.eval_biasbios_instruction \
--task json \
--apply_pasta \
--emphasized_text instruct \
--alpha 0.01 \
--scale_position exclude \
--pasta_head_config <head_config_path> \
--model huggyllama/llama-7b \
--prompt_idx 0 \
--batch_size 16 \
--max_new_tokens 128 \
--experiment_name llama_evaluation \
--device cuda 

Pronouns Changing

python -m scripts.eval_biasbios_instruction \
--task pronchange \
--apply_pasta \
--emphasized_text instruct \
--alpha 0.01 \
--scale_position exclude \
--pasta_head_config <head_config_path> \
--prompt_idx 0 \
--model huggyllama/llama-7b \
--max_new_tokens 128 \
--batch_size 16 \
--experiment_name llama_evaluation \
--device cuda 

Bias in Bios

python -m scripts.eval_bias_gen \
--model huggyllama/llama-7b \
--apply_pasta \
--alpha 0.01 \
--scale_position exclude \
--pasta_head_config <head_config_path> \
--max_length 256 \
--batch_size 16 \
--experiment_name llama_evaluation \
--device cuda 

CounterFact

python -m scripts.eval_fact_gen \
--model huggyllama/llama-7b \
--apply_pasta \
--alpha 0.01 \
--scale_position exclude \
--pasta_head_config <head_config_path> \
--add_unmediated_fact True \
--benchmarks efficacy paraphrase generation \
--experiment_name llama_evaluation

Contact

Please contact us or post an issue if you have any questions:

• Qingru Zhang ([e-mail])
• Chandan Singh ([e-mail])
• Liyuan Liu ([e-mail])
• Xiaodong Liu ([e-mail])

@misc{zhang2023tell,
    title={Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs},
    author={Qingru Zhang and Chandan Singh and Liyuan Liu and Xiaodong Liu and Bin Yu and Jianfeng Gao and Tuo Zhao},
    year={2023},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}