pythia-1B-feature-extraction-wikitext-103

Pythia-1B Feature-Based Transfer Learning on WikiText-103

This repository contains a feature-based transfer learning experiment using the EleutherAI/pythia-1b model on the WikiText-103 dataset.
The base model was frozen, and a lightweight trainable classification head was added for causal language modeling.

---

Model Description

• Base Model: EleutherAI/pythia-1b
• Training Paradigm: Feature-based transfer learning (frozen base + new lightweight head)
• Task: Causal Language Modeling
• Dataset: WikiText-103 (raw v1)

The base model (gpt_neox) was frozen to retain pretrained knowledge. A new head (2-layer feedforward with ReLU and dropout) was trained on top of the hidden states for efficient adaptation.

---

Training Setup

• Framework: Transformers + PyTorch
• GPU: Multi-GPU (CUDA enabled)
• Batch size: 8 (gradient accumulation: 2)
• Sequence length (block size): 1024
• Optimizer: AdamW
• Learning rate: 2e-4 with cosine decay
• Epochs: 10
• Mixed Precision: FP16
• Callbacks: Early stopping, custom metric logging

---

Results

Final Training Metrics

• Training Loss: 2.6275
• Final Step Loss: 2.4289
• Gradient Norm: 0.3317
• Learning Rate at End: 1.55e-06

Evaluation Metrics (Epoch 10)

• Evaluation Loss: 2.5432
• Evaluation Perplexity: 12.72
• Evaluation Runtime: 1.6039s
• Samples per Second: 150.26
• Steps per Second: 4.99

---

Usage

from transformers import AutoTokenizer
import torch
from model import FrozenPythiaWithNewHead 

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("./pythia-wikitext-feature")

# Load model
model = FrozenPythiaWithNewHead.from_pretrained("./pythia-wikitext-feature")
model.eval()

# Example
input_text = "The history of natural language processing"
inputs = tokenizer(input_text, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs["logits"]
    next_token_id = torch.argmax(logits[:, -1, :], dim=-1)
    print("Next token:", tokenizer.decode(next_token_id))