Selective Prompt Anchoring

Selective Prompt Anchoring (SPA) is a model-agnostic algorithm designed for large language models (LLMs) that provides fine-grained control over text generation.

This is the official repo for our ICML 2025 work.

📄 Paper: Selective Prompt Anchoring for Code Generation

🌐 Live Demo: Try me 👈🏻

⏯️ Demo Video (Code Generation): Watch me 👈🏻

🖼️ Poster

Why use SPA? 🤔

When humans talk with one another, nuanced emphasis and fine-grained implications are often conveyed through different volumes, tones, or pauses. However, there is a gap in conveying such information through text-based communication with AI. Moreover, AI may attend to the wrong parts of the input. We need an efficient mechanism to correct its attention, without reference challenges (i.e., it is hard to reference when the context is repetitive or scattered), and without increasing the context (i.e., simply repeating may not only make the prompt weird but also risk exceeding the context window size).

✨ SPA enables users to assign importance, emphasis, or weights to specific parts of input text when prompting LLMs. SPA brings this capability to text-based AI communication by allowing users to anchor (the name is inspired by anchoring effect in psychology) certain words or phrases in the prompt, causing the model to pay more/less/reversed attention to them during generation. With SPA, users can flexibly steer LLMs' attention through the Hugging Face API.

👉 Here are some interesting examples

Example 1

Example 2

Example 3

Note: While we currently work on text-to-text generation and evaluating on code generation in our paper, the underlying idea can be potentially applied to other tasks (e.g., classification) or with other modalities (e.g., image).

How SPA Works 💡

At each step, SPA generates two logit distributions in parallel, based on:

1. The user prompt (+ alreadly generated tokens)
2. The prompt with anchored tokens masked (+ already generated tokens)

Then, SPA compares the two logit distributions and adjusts the final probabilities, where the influence of anchored text is increased. The following figure demonstrates the high-level idea:

Installation 📦

From PyPI

(Recommended) Install directly from PyPI using pip:

pip install anchoring

# pip install accelerate (Sometimes need to do this after the first command)

From GitHub

git clone https://github.com/your-username/selective-prompt-anchoring.git
cd selective-prompt-anchoring
pip install -e .

From HuggingFace

pip install huggingface_hub
pip install git+https://github.com/magic-YuanTian/Selective-Prompt-Anchoring.git

Quick Start ⚡

With pipeline API:

from transformers import pipeline
import anchoring

pipe = pipeline(
    "selective-prompt-anchoring",
    model="meta-llama/Llama-3.1-8B-Instruct",
)

output = pipe("How is the weather today?", anchors=['today'])
print(output["generated_text"])

A bit more settings:

from transformers import pipeline
import anchoring  # The pipeline is automatically registered on import

# Create pipeline
spa_pipe = pipeline(
    "selective-prompt-anchoring",
    model="meta-llama/Llama-3.1-8B-Instruct",
    anchoring_strength=3.0,
    modulated_by_prob=True,
    use_attention_mask=True,
    device_map="auto"
)

# Simple text prompt with global anchors
prompt = "How is the weather today?"
global_anchors = ['today']

output = spa_pipe(prompt, anchors=global_anchors, max_new_tokens=1024)
print(output["generated_text"])

You can also stream the generation

SPA supports streaming for real-time generation:

# Get streaming output
for token in spa_pipe(prompt, anchors=global_anchors, max_new_tokens=1024, stream=True):
    print(token, end="", flush=True)
print()

Batch processing

# Define a list of prompts
prompts = ["What's the weather <anchor>today</anchor>?", "What's the weather <anchor>tomorrow</anchor>?"]

# Or with chat format
prompts = [
    [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What's the weather <anchor>today</anchor>?"}
    ],
    [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What's the weather <anchor>tomorrow</anchor>?"}
    ]
]

# Process all prompts
outputs = spa_pipe(prompts, anchors=['weather'], max_new_tokens=1024)
for output in outputs:
    print(output["generated_text"])

Direct Usage with model.generate()

(Recommended for developers)

This option is more compatible with Hugging Face API and potentially supports more parameters and models.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from anchoring import SPALogitsProcessor, spa_tokenize

# Load model and tokenizer
model_name = "meta-llama/Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# Define anchors and prompt
global_anchors = ['today']
prompt = "How is the weather today?"

# Tokenize with SPA
main_inputs, aux_inputs, mask_token = spa_tokenize(
    prompt_with_anchors=prompt,
    global_anchors=global_anchors,
    tokenizer=tokenizer,
    device=model.device
)

# Create SPA logits processor
spa_processor = SPALogitsProcessor(
    aux_model=model, 
    aux_input_ids=aux_inputs, 
    strength=3.0,
    modulated_by_prob=False,
    use_attention_mask=True,
    mask_token=mask_token,
    tokenizer=tokenizer
)

# Generate text with SPA
output_sequences = model.generate(
    input_ids=main_inputs,
    attention_mask=torch.ones_like(main_inputs),
    logits_processor=[spa_processor],
    max_new_tokens=1024,
    do_sample=False,
)

# Decode and print
generated_text = tokenizer.decode(output_sequences[0], skip_special_tokens=True)
print(generated_text)

Input Formats

Our code supports multiple input formats, allowing developers to conveniently represent anchors in prompts or messages. Developers can use inline paired tags, <anchor> </anchor>, or a global anchor list to denote anchored text. They can also work with chat messages in a list, following the OpenAI API standard, or simply use a prompt string.

1️⃣ String with Global Anchors

prompt = "How is the weather today?"
global_anchors = ['today']

2️⃣ String with Inline Anchors

prompt = "What's the weather <anchor>today</anchor>? Think <anchor>step by step</anchor>."

3️⃣ Chat Messages with Message-Level Anchors

prompt = [
    {
        "role": "system", 
        "content": "You are a helpful assistant.", 
        "anchors": ["You", "assistant"]  
    },
    {
        "role": "user",
        "content": "What's the weather today?", 
        "anchors": ["today"]
    },
]

4️⃣ Chat Messages with Inline Anchors

prompt = [
    {
        "role": "system", 
        "content": "You are a helpful assistant."
    },
    {
        "role": "user", 
        "content": "What's the weather <anchor>today</anchor>?"
    },
]

👉 Two example usage scripts are included under the example folder!

Hyper-parameters

SPA-Specific Parameters

• strength (default: 1.4): Controls the influence of anchored text.

  • 1.0: No effect (normal generation)
  • 0.0: Completely ignore anchored text
  • >1.0: Emphasize anchored text (higher values = stronger emphasis)
  • <0.0: Reverse the influence of the anchored text (negative values = stronger reversed influence)

• modulated_by_prob (default: True): When True, the anchoring strength is modulated by token probability.

  • Enable for more stable results, especially with higher anchoring strengths
  • Disable for more precise control at lower strengths
  • Hint: Set this as False to reproduce the paper settings

• use_attention_mask (default: True): When True, uses attention masking for anchor tokens, enhancing the effectiveness of anchoring.

  • If you observe unexpected behavior, set this as False
  • Hint: Set this as False to reproduce the paper setting

Standard Generation Parameters

SPA supports all standard Hugging Face generation parameters, such as:

• max_new_tokens: Maximum number of tokens to generate
• do_sample: Whether to use sampling for generation
• temperature: Controls randomness (higher = more random)
• top_p: Top-p sampling parameter (nucleus sampling)
• top_k: Top-k sampling parameter
• min_new_tokens: Minimum number of tokens to generate

For more parameters, please check the official Huggingface Transformers' generation documentation.

Practical Hyper-parameter Settings

1. strength (Anchoring Strength):
   • When you want to increase the model's attention/text emphasis
     • If modulated_by_prob = True, you can give a relatively high value of anchoring strength (e.g., 20).
     • If modulated_by_prob = False, we recommend a value less than 2.
     • If you are pursuing an optimal value, you can easily tune this value through grid search on your benchmark. Our experiment demonstrates that this value follows a simple pattern (as value increases, performance first improves, then declines), and it is easy to tune by dozens of examples.
   • For reducing (0 < anchoring_strength < 1) or reversing (anchoring_strength < 0), please set the value based on your concrete needs.

Note: In our ICML'25 paper, we place less emphasis on modulated_by_prob and focus more on the tuning method to demonstrate the impact of this strength value. For further tuning details, please refer to our paper.

2. modulated_by_prob (Weight influence by token probabilities): We recommend setting modulated_by_prob=True for stable results. Set it as False if you aim for precise control or have other development needs.

3. use_attention_mask (whether to use attention mask or just special token masking): Set True by default for more reliable performance, unless you detect any performance issue, you can set it as False, SPA supports a backup masking strategy by special tokens.

Compatibility

SPA is a model-agnostic algorithm. Our implementation inherits the Huggingface Transformers generation API. It should work for ANY LLM from Huggingface model collections. Please follow the corresponding model documentation for detailed instructions.

If you have any questions or need support for a specific model, please don't hesitate to submit an issue. We will respond shortly. 😁

Citation

If you find SPA useful, you could cite by

@inproceedings{tian2025spa,
  title = {Selective Prompt Anchoring for Code Generation},
  author = {Tian, Yuan and Zhang, Tianyi},
  booktitle = {Proceedings of the 42nd International Conference on Machine Learning (ICML)},
  year = {2025},
  url = {https://arxiv.org/abs/2408.09121},
  eprint = {2408.09121},
  archivePrefix = {arXiv}
}

Contact

• Email: tian211@purdue.edu
• Website: yuan-tian.com