A first experiment to test and convert clip-vit-base-patch32 into a geometric model by using only a classification head.
Below is GPT 5's auto-generated dictation based on the notebook. I'll include the full notebook in a moment here.
The answer is... maybe. More research required.
I used the 32 dim geometric vocab; as it seemed to be the weakest with flow-match euler-discreet to test the hypothesis that a small dimensional geometry could in fact be used in substitution of a high-geometric variation.
The output head is incredibly small unlike my first impression. This one happened to have the pair packed, but the updated notebook will show the separated head is in fact less than 100 kb.
I believe the clip-vit variations have more utility overall so I wanted to ensure a fair target was assessed.
One-vector image embeddings (HF CLIP) + pentachora vocabulary anchors → cosine-similarity classifier for CIFAR-100.
This repo hosts the trained crystal classification head (+ run configs/metrics) built in Notebook 6.
Note: If you only want to ship the head, you can also include a stripped crystal_head.safetensors (head-only state_dict). The snippets below handle either format.
-
Load CLIP vision (frozen) and processor
HF_CLIP_ID = "openai/clip-vit-base-patch32"
Processor = AutoImageProcessor.from_pretrained(HF_CLIP_ID)
Encoder = CLIPVisionModelWithProjection.from_pretrained(HF_CLIP_ID).eval().to("cuda")
-
Build the crystal head (same shape as training)
image_dim = Encoder.config.projection_dim # 512
crystal_dim = 512 # vocab repo uses 512D anchors
sym_dim = 128 # crystal_dims from CONFIG
temperature = 0.07 # from CONFIG
class CrystalHead(torch.nn.Module):
def init(self, De, Dv, Dsym, T):
super().init()
self.proj_img = torch.nn.Linear(De, Dsym, bias=True)
self.proj_anc = torch.nn.Linear(Dv, Dsym, bias=False)
self.T = T
self.register_buffer("anchors_vocab", torch.empty(0, Dv), persistent=False)
def set_anchors(self, anchors): # [C, Dv]
self.anchors_vocab = anchors.contiguous()
def forward(self, image_embeds): # [B, De] (L2 ok)
z = torch.nn.functional.normalize(self.proj_img(image_embeds), dim=-1)
a = torch.nn.functional.normalize(self.proj_anc(self.anchors_vocab), dim=-1)
return (z @ a.T) / max(1e-8, self.T) # [B, C]
head = CrystalHead(De=image_dim, Dv=crystal_dim, Dsym=sym_dim, T=temperature).to("cuda")
-
Load weights (handles prefixed multi-module .safetensors)
state = safetensors.torch.load_file("<run_name>_best.safetensors")
head_state = {k.split("head::",1)[1]: v for k,v in state.items() if k.startswith("head::")}
head.load_state_dict(head_state, strict=True)
-
Prepare anchors from your vocabulary (same order as training)
You likely already exported anchors or can rebuild them exactly as in Notebook 6.
anchors: torch.Tensor of shape [100, 512]
head.set_anchors(anchors.to("cuda"))
-
Inference on a batch of images (PIL or ndarray)
imgs = [PIL.Image.open("example_0.png").convert("RGB"), PIL.Image.open("example_1.png").convert("RGB")]
batch = Processor(images=imgs, return_tensors="pt").to("cuda")
with torch.no_grad():
out = Encoder(pixel_values=batch["pixel_values"], return_dict=True)
z = torch.nn.functional.normalize(out.image_embeds, dim=-1) # [B, 512]
logits = head(z) # [B, 100]
pred = logits.argmax(dim=-1).tolist()
print("pred:", pred)
Note: The head expects the same class order used at training time. Save and ship class_names.json (CIFAR-100 labels) and the exact anchors_vocab.pt you used (or rebuild deterministically with the vocab + synth step).
Replace with your final numbers after the run completes.
