Add V-JEPA 2 archetypal transcoder (headline EV 0.480 + random-init control, model card, code)
3b22208 verified | """Minimal, self-contained loader for the V-JEPA 2 archetypal sparse transcoder. | |
| Requires: torch, safetensors, and modeling_archetypal_transcoder.py (this repo). | |
| """ | |
| import torch | |
| from safetensors.torch import load_file | |
| from modeling_archetypal_transcoder import ArchetypalTranscoder, ArchetypalConfig | |
| def load_transcoder(weights="model.safetensors", device="cpu"): | |
| sd = load_file(weights) | |
| cfg = ArchetypalConfig( | |
| d_in=768, d_hidden=12288, n_pool=32000, top_k=512, | |
| init_eye=True, delta=0.5, encoder_topk=None, | |
| ) | |
| # data_pool (K-means centroids) is bundled in the weights; pass it to the ctor. | |
| model = ArchetypalTranscoder(cfg, data_pool=sd["data_pool"]).to(device) | |
| model.load_state_dict(sd) | |
| model.eval() | |
| return model | |
| if __name__ == "__main__": | |
| m = load_transcoder() | |
| # x = V-JEPA 2.1 block-10 FFN *input* activations, shape [N, 768], fp32. | |
| x = torch.randn(4, 768) | |
| x_hat, z = m(x) # x_hat: predicted FFN output [N,768]; z: sparse codes [N,12288] | |
| print("x_hat", tuple(x_hat.shape), "| z", tuple(z.shape), "| L0", (z>0).float().sum(-1).mean().item()) | |