jahnavisharma/llm-representation-alignment

LLM Representation Alignment

Overview

This repository contains the initial experiments for investigating representation alignment and model stitching in Large Language Models (LLMs).

Goal

Our objective is to determine whether hidden-state representations from different LLMs can be aligned and whether representation alignment can help explain hallucinations and factual errors.

Models

• TinyLlama/TinyLlama-1.1B-Chat-v1.0
• Qwen/Qwen2.5-0.5B

Dataset

• TruthfulQA
• Small custom factual QA dataset (for initial testing)

Current Progress

• Set up local experimentation environment using Hugging Face Transformers and Ollama.
• Loaded and evaluated TinyLlama and Qwen.
• Switched from a small custom dataset to TruthfulQA.
• Extracted hidden states from all transformer layers.
• Compared hidden-state representations.
• Observed that:
  • TinyLlama hidden size = 2048
  • Qwen hidden size = 896

Project Structure

• evaluate.py – Runs a small factual QA evaluation.
• benchmark_test.py – Loads and explores the TruthfulQA benchmark.
• hidden_states.py – Extracts hidden states from all transformer layers.
• compare_models.py – Compares hidden-state representations across models.
• factual_qa.json – Small custom dataset used for initial testing.

Future Work

• Evaluate on a larger TruthfulQA subset.
• Compare representations across different layers.
• Project hidden states into a common latent space.
• Investigate representation stitching between models.
• Study the relationship between representation similarity and hallucinations.

Reproducibility

The complete source code, setup instructions, and scripts needed to reproduce these experiments are available on GitHub:

GitHub Repository: https://github.com/jsharma0110/llm-representation-alignment