Skip to content

README.md

Latest commit

 

History

History
270 lines (198 loc) · 9.29 KB

README.md

File metadata and controls

270 lines (198 loc) · 9.29 KB

Task Vector Interpretation and Analysis

A project exploring task vectors in large language models (LLMs) through fine-tuning, extraction, and analysis across multiple NLP tasks and model architectures.

📋 Overview

This repository implements and analyzes task vectors - the difference between fine-tuned and pretrained model parameters that capture task-specific knowledge. The project demonstrates how these vectors can be extracted, stored, and applied to models for efficient transfer learning.

Key Features

  • 🔍 Task Vector Extraction: Extract task-specific knowledge from fine-tuned models
  • 🔄 Model Application: Apply task vectors to pretrained models for instant specialization
  • 📊 Multi-Task Analysis: Support for MNLI (Natural Language Inference) and SST-2 (Sentiment Analysis)
  • 🤖 Multi-Model Support: Compatible with Llama 3.1, Mistral, and Gemma architectures
  • ⚡ Efficient Fine-tuning: Uses LoRA (Low-Rank Adaptation) for parameter-efficient training
  • 💾 Portable Vectors: Save and load task vectors as lightweight .pth files

🏗️ Project Structure

Task_vector_Interpretation_and_Analysis/
├── MNLI/                          # Multi-Genre Natural Language Inference
│   ├── gemma/
│   │   └── gemma2-9b-mnli.ipynb
│   ├── mistral/
│   │   └── [mistral MNLI notebooks]
│   └── llama/
│       ├── llama3-1-8b-mnli.ipynb
│       └── mnli_llama3.1_8B_task_vector.pth
├── SST-2/                         # Stanford Sentiment Treebank
│   ├── gemma/
│   │   └── [gemma SST-2 notebooks]
│   ├── mistral/
│   │   └── [mistral SST-2 notebooks]
│   └── llama/
│       ├── llama3-1-8b-sst-2.ipynb
│       └── sst2_llama3.1_8B_task_vector.pth
└── Load Task vector from file and merge with model/
    └── load-task-vector-from-pth-file.ipynb

🚀 Quick Start

Prerequisites

pip install torch torchvision torchaudio xformers --index-url https://download.pytorch.org/whl/cu121
pip install unsloth transformers datasets trl

1. Fine-tune a Model and Extract Task Vector

from unsloth import FastLanguageModel
import torch

# Load pretrained model
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/Meta-Llama-3.1-8B",
    max_seq_length=2048,
    load_in_4bit=True,
)

# Apply LoRA for efficient fine-tuning
model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", 
                   "gate_proj", "up_proj", "down_proj"],
    lora_alpha=16,
    use_gradient_checkpointing="unsloth",
)

# Fine-tune on your task (MNLI, SST-2, etc.)
# ... training code ...

# Extract task vector
task_vector = TaskVector(pretrained_model, fine_tuned_model)
torch.save(task_vector.vector, "my_task_vector.pth")

2. Load and Apply Task Vector

# Load saved task vector
task_vector_data = torch.load("sst2_llama3.1_8B_task_vector.pth")
task_vector = TaskVector(vector=task_vector_data)

# Apply to pretrained model
specialized_model = task_vector.apply_to(pretrained_model)

# Use for inference
inputs = tokenizer("This movie is amazing!", return_tensors="pt")
outputs = specialized_model.generate(**inputs)

📚 Supported Tasks

MNLI (Multi-Genre Natural Language Inference)

  • Task: Determine relationship between premise and hypothesis
  • Labels: entailment, neutral, contradiction
  • Dataset: GLUE MNLI benchmark

SST-2 (Stanford Sentiment Treebank)

  • Task: Binary sentiment classification
  • Labels: positive, negative
  • Dataset: Movie review sentiment analysis

🤖 Supported Models

Model Size Architecture Status
Llama 3.1 8B Transformer ✅ Fully Supported
Mistral 7B Transformer ✅ Fully Supported
Gemma 2 9B Transformer ✅ Fully Supported

🧠 Task Vector Methodology

What are Task Vectors?

Task vectors represent the difference between a fine-tuned model and its pretrained base:

Task Vector = Fine-tuned Model Parameters - Pretrained Model Parameters

Key Advantages

  1. Modularity: Task-specific knowledge stored separately from base model
  2. Efficiency: Vectors are much smaller than full model weights
  3. Composability: Multiple task vectors can potentially be combined
  4. Portability: Easy to share and apply across different instances

Implementation Details

  • LoRA Fine-tuning: Parameter-efficient training using Low-Rank Adaptation
  • Chunked Application: Memory-efficient vector application for large models
  • GPU/CPU Flexibility: Optimized for both GPU and CPU inference
  • Platform: All notebooks developed and tested on Kaggle with GPU T4 x2 using Unsloth for accelerated training

📊 Results and Analysis

The project demonstrates successful task vector extraction and application across:

  • ✅ MNLI Classification: Accurate entailment/contradiction detection
  • ✅ Sentiment Analysis: Reliable positive/negative classification
  • ✅ Cross-Model Transfer: Vectors work across different model architectures
  • ✅ Memory Efficiency: Significant storage savings compared to full model weights

🔬 Research Applications

This codebase enables research in:

  • Task Vector Arithmetic: Combining and manipulating task knowledge
  • Transfer Learning: Efficient knowledge transfer between tasks
  • Model Interpretability: Understanding what models learn for specific tasks
  • Few-Shot Learning: Rapid adaptation to new tasks
  • Model Compression: Storing task-specific knowledge efficiently

📁 File Descriptions

Notebooks

  • *-mnli.ipynb: Complete pipeline for MNLI task vector creation
  • *-sst-2.ipynb: Complete pipeline for SST-2 task vector creation
  • load-task-vector-from-pth-file.ipynb: Demonstration of loading and applying saved vectors

Task Vector Files

  • *.pth: Serialized task vectors (PyTorch format)
  • Typically 2-3GB per vector, much smaller than full model weights

📥 Downloading Pre-trained Task Vectors

Due to GitHub's file size limitations, the large task vector files (.pth) are not included in this repository. You can:

  1. Generate your own: Follow the notebooks to create task vectors from scratch
  2. Download pre-trained vectors: [Contact for access to pre-trained task vectors]
  3. Use Git LFS: Clone with git lfs clone if LFS is enabled

Available pre-trained task vectors:

  • sst2_llama3.1_8B_task_vector.pth (2.5GB) - Llama 3.1 8B for SST-2
  • mnli_llama3.1_8B_task_vector.pth (2.5GB) - Llama 3.1 8B for MNLI
  • sst2_gemma2_9B_task_vector.pth (3.7GB) - Gemma 2 9B for SST-2
  • sst2_mistral_instruct_v0.3_7B_task_vector.pth (675MB) - Mistral 7B for SST-2
  • mnli_mistral_nemo_12B_task_vector.pth (675MB) - Mistral Nemo 12B for MNLI

🛠️ Technical Requirements

Development Environment

  • Platform: Kaggle Notebooks (recommended)
  • GPU: NVIDIA Tesla T4 x2 (used in development)
  • Framework: Unsloth for efficient fine-tuning

System Requirements

  • GPU: NVIDIA GPU with CUDA support (Tesla T4 or better)
  • Memory: 16GB+ GPU memory recommended for 8B+ models
  • Storage: ~10GB per model + task vectors
  • Python: 3.8+ with PyTorch 2.0+

Note on Environment

All Jupyter notebooks (*.ipynb) in this repository were developed and executed on Kaggle using GPU T4 x2 with the Unsloth library for optimized training.

📖 Usage Examples

Sentiment Analysis with Task Vector

# Load SST-2 task vector
sst2_vector = torch.load("sst2_llama3.1_8B_task_vector.pth")
task_vector = TaskVector(vector=sst2_vector)

# Apply to model
model = task_vector.apply_to(base_model)

# Classify sentiment
prompt = """Below is an instruction that describes a task...
### Input:
This movie is fantastic!
### Response:
"""
result = model.generate(tokenizer.encode(prompt))
# Output: "positive"

Natural Language Inference

# Load MNLI task vector
mnli_vector = torch.load("mnli_llama3.1_8B_task_vector.pth")
task_vector = TaskVector(vector=mnli_vector)

# Apply to model
model = task_vector.apply_to(base_model)

# Perform inference
premise = "A person is reading a book"
hypothesis = "Someone is learning"
# Output: "entailment"

🤝 Contributing

Contributions are welcome! Areas for improvement:

  • Additional task implementations (summarization, QA, etc.)
  • More model architectures (GPT, T5, etc.)
  • Task vector arithmetic and composition
  • Evaluation metrics and benchmarking
  • Documentation and tutorials

📜 License

This project is open-source and available under the MIT License.

🙏 Acknowledgments

  • Unsloth: For efficient model fine-tuning capabilities
  • Hugging Face: For model hosting and datasets
  • Task Vector Research: Based on emerging research in modular AI systems

📚 References

Note: This project is designed for research and educational purposes. Task vectors represent an active area of AI research with ongoing developments in methodology and applications.