In-Context Learning Emergence › Studying ICL Emergence Empirically

import numpy as np
import torch


def measure_icl_ability(model_name, task_examples, test_cases, device="cpu"):
    """
    Measure a model's in-context learning ability on a simple task.

    Returns accuracy improvement from adding demonstrations.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
    model.eval()

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    def get_prediction(prompt, choices):
        """Get model's preferred choice."""
        choice_probs = []

        for choice in choices:
            full_text = prompt + choice
            inputs = tokenizer(full_text, return_tensors="pt").to(device)

            with torch.no_grad():
                outputs = model(**inputs)
                # Get probability of last token
                logits = outputs.logits[0, -1, :]
                choice_token_id = tokenizer.encode(
                    choice, add_special_tokens=False
                )[0]
                prob = torch.softmax(logits, dim=-1)[choice_token_id].item()
                choice_probs.append(prob)

        return choices[np.argmax(choice_probs)]

    # Format demonstrations
    demo_text = "\n".join(
        [
            f"Input: {ex['input']}\nOutput: {ex['output']}"
            for ex in task_examples
        ]
    )

    zero_shot_correct = 0
    few_shot_correct = 0

    for test in test_cases:
        choices = test["choices"]
        correct = test["correct"]

        # Zero-shot
        zero_shot_prompt = f"Input: {test['input']}\nOutput: "
        zero_shot_pred = get_prediction(zero_shot_prompt, choices)
        if zero_shot_pred == correct:
            zero_shot_correct += 1

        # Few-shot
        few_shot_prompt = demo_text + f"\nInput: {test['input']}\nOutput: "
        few_shot_pred = get_prediction(few_shot_prompt, choices)
        if few_shot_pred == correct:
            few_shot_correct += 1

    n = len(test_cases)
    return {
        "zero_shot_acc": zero_shot_correct / n,
        "few_shot_acc": few_shot_correct / n,
        "icl_gain": (few_shot_correct - zero_shot_correct) / n,
    }

# Define a simple sentiment classification task
task_examples = [
    {"input": "This movie was fantastic", "output": "positive"},
    {"input": "I hated every minute", "output": "negative"},
    {"input": "Best experience ever", "output": "positive"},
    {"input": "Terrible waste of time", "output": "negative"},
]

test_cases = [
    {
        "input": "Absolutely wonderful",
        "choices": ["positive", "negative"],
        "correct": "positive",
    },
    {
        "input": "Disappointing result",
        "choices": ["positive", "negative"],
        "correct": "negative",
    },
    {
        "input": "I loved it",
        "choices": ["positive", "negative"],
        "correct": "positive",
    },
    {
        "input": "Would not recommend",
        "choices": ["positive", "negative"],
        "correct": "negative",
    },
]

import numpy as np
import torch


def analyze_icl_representations(model, tokenizer, demonstrations, query):
    """
    Analyze how representations change through the layers during ICL.

    This probes whether task information accumulates in specific layers.
    """
    # Prepare input
    demo_text = "\n".join([f"Q: {d['q']}\nA: {d['a']}" for d in demonstrations])
    full_prompt = demo_text + f"\nQ: {query}\nA:"

    inputs = tokenizer(full_prompt, return_tensors="pt")

    # Get all hidden states
    with torch.no_grad():
        outputs = model(**inputs, output_hidden_states=True)

    hidden_states = (
        outputs.hidden_states
    )  # Tuple of (batch, seq_len, hidden_dim)

    # Analyze the final position (where prediction happens)
    final_position_states = [h[0, -1, :].numpy() for h in hidden_states]

    # Compute change in representation across layers
    layer_changes = []
    for i in range(1, len(final_position_states)):
        change = np.linalg.norm(
            final_position_states[i] - final_position_states[i - 1]
        )
        layer_changes.append(change)

    return {
        "final_states": final_position_states,
        "layer_changes": layer_changes,
    }