Qwen2.5-1.5B-Instruct NLA L18 Activation Reconstructor

The activation reconstructor of a Natural Language Autoencoder trained on the layer-18 residual stream of Qwen/Qwen2.5-1.5B-Instruct. See https://transformer-circuits.pub/2026/nla/, and https://github.com/kitft/natural_language_autoencoders

Activation verbalizer: dormantx/Qwen2.5-1.5B-Instruct-NLA-L18-av.

Usage

import torch, torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModel, AutoTokenizer

repo = "dormantx/Qwen2.5-1.5B-Instruct-NLA-L18-ar"
tok = AutoTokenizer.from_pretrained(repo)
ar = AutoModel.from_pretrained(repo, torch_dtype=torch.bfloat16).cuda().eval()
head = load_file(hf_hub_download(repo, "value_head.safetensors"))
W, b = head["weight"].cuda(), head["bias"].cuda()

if tok.pad_token is None:
    tok.pad_token = tok.eos_token
enc = tok(["the outcome of a football game"], return_tensors="pt", padding=True).to("cuda")
with torch.no_grad():
    h = ar(**enc, output_hidden_states=True).hidden_states[18]      # [B, T, 1536]
    last = enc["attention_mask"].sum(1) - 1
    pooled = h[torch.arange(h.shape[0]), last].to(torch.float32)    # last non-pad token
    pred = pooled @ W.T.float() + b.float()                         # reconstructed direction
    pred = F.normalize(pred, dim=-1)