Feature Request: Multi-Task Prompt Tuning Support

[hwilner]

Summary

I would like to propose adding support for Multi-Task Prompt Tuning to the prompt-tuning library. This technique enables efficient learning across multiple related tasks by decomposing prompts into shared and task-specific components.

Motivation

Multi-task prompt tuning facilitates knowledge transfer between tasks while maintaining task-specific adaptations. This is particularly valuable when working with multiple related tasks, as it can improve performance and parameter efficiency compared to training separate prompts for each task.

Proposed Implementation

The implementation would include:

1. Shared and task-specific components: Decompose prompts into parts that are shared across tasks and parts that are task-specific
2. Multiple composition methods: Support concatenation, addition, weighted combination, and gated composition
3. Hierarchical organization: Enable grouping of related tasks for better knowledge sharing
4. Adaptive composition: Use attention mechanisms to dynamically combine shared and task-specific prompts

Key Features

• Shared prompts across multiple tasks
• Task-specific prompt components
• Flexible composition strategies
• Hierarchical task grouping
• Integration with existing T5X/Flaxformer infrastructure

Reference

Wang et al. (2022). "Multitask Prompt Tuning Enables Parameter-Efficient Transfer Learning." arXiv:2303.02861

Additional Context

I have implemented a prototype of this technique that follows the library's design patterns and coding standards. The implementation is available in my fork at https://github.com/hwilner/prompt-tuning

Would the maintainers be interested in this enhancement? I'm happy to discuss the design and implementation details further.