Expose patch tokens alongside the CLS embedding

The backbone already computed all 789 tokens, so returning the 784 patch ones
costs no extra inference. Registers stay dropped; they are artifact sinks, not
features. Patches are unnormalized because dense heads want the magnitude.

Dropping CLS from the artifact name since it no longer describes the outputs.
The weights are byte-identical, so LFS reuses the object already pushed.

README.md

@@ -14,8 +14,9 @@ tags:
 
 # DINOv3 Core ML
 
-A Core ML conversion of Meta's DINOv3 ViT-B/16 image encoder. It returns one
-L2-normalized, 768-value CLS embedding per image for image similarity and retrieval.
+A Core ML conversion of Meta's DINOv3 ViT-B/16 image encoder. It returns an
+L2-normalized, 768-value CLS embedding for image similarity and retrieval, plus the
+784 raw patch tokens for dense tasks such as segmentation, depth, and correspondence.
 This is a format conversion of pretrained weights; no additional training was performed.
 
 ## Model
@@ -24,18 +25,24 @@ This is a format conversion of pretrained weights; no additional training was pe
 |---|---|
 | Base model | `facebook/dinov3-vitb16-pretrain-lvd1689m` |
 | timm implementation | `vit_base_patch16_dinov3.lvd1689m` |
-| Artifact | `models/DINOv3ViTB16CLS-FP32-448.mlpackage` |
+| Artifact | `models/DINOv3ViTB16-FP32-448.mlpackage` |
 | Format | Core ML ML Program, FP32 |
 | Target | macOS 14 or newer |
 | Input | `image`: 448×448 RGB, pixel values 0–255 |
 | Output | `embedding`: float32, shape `[1, 768]`, unit L2 norm |
+| Output | `patch_embeddings`: float32, shape `[1, 784, 768]`, unnormalized |
 
 Correct EXIF orientation, convert to RGB, and resize to 448×448 before inference.
 The examples use a square resize, which can distort non-square images. Use the same
 preprocessing for all images being compared. Pixel scaling and ImageNet normalization
-are inside the model; do not apply them again. The output uses the CLS token only,
-not patch or register tokens. Similarity is the dot product of two output vectors.
-It does not provide text embeddings, captions, or geographic coordinates.
+are inside the model; do not apply them again. Similarity is the dot product of two
+`embedding` vectors. It does not provide text embeddings, captions, or coordinates.
+
+`patch_embeddings` is a 28x28 grid of 768-value tokens flattened in row-major order,
+one per 16x16 input patch. Unlike `embedding` these are not normalized, because dense
+heads generally want the magnitude; normalize per token yourself for cosine. The four
+register tokens are dropped: they exist to absorb high-norm artifacts that would
+otherwise pollute the patch tokens, and are not useful as features.
 
 ## Run
 
@@ -56,7 +63,7 @@ For native Swift, provide an already resized, orientation-corrected 448×448 ima
 ```sh
 mkdir -p build
 swiftc -O examples/encode.swift -o build/encode
-build/encode models/DINOv3ViTB16CLS-FP32-448.mlpackage image-448.png
+build/encode models/DINOv3ViTB16-FP32-448.mlpackage image-448.png
 ```
 
 The example compiles the package at runtime. An app can instead add the package to
@@ -66,21 +73,21 @@ resource once for repeated predictions.
 ## Convert
 
 ```sh
-uv run python convert.py --output build/DINOv3ViTB16CLS-FP32-448.mlpackage
+uv run python convert.py --output build/DINOv3ViTB16-FP32-448.mlpackage
 ```
 
 The converter downloads pretrained weights through timm/Hugging Face. If the upstream
 checkpoint requires access, accept its terms and authenticate with Hugging Face first.
 It exports through `torch.export`, converts to Core ML, and compares both runtimes
-on a synthetic RGB gradient. A cosine below 0.999 fails validation.
+on a synthetic RGB gradient. A cosine below 0.999 on either output fails validation.
 
 `--size` changes the square input resolution; use a positive multiple of 16.
 `--precision fp16` is experimental and must pass its own validation.
 Existing output packages are protected unless `--force` is supplied.
 
 The copied FP32 artifact's original validation result is in
-`models/DINOv3ViTB16CLS-FP32-448.validation.json`: cosine approximately 1.0
-and output norm approximately 1.0. Floating-point rounding can put cosine slightly
+`models/DINOv3ViTB16-FP32-448.validation.json`: cosine approximately 1.0 for both
+outputs, and CLS norm approximately 1.0. Floating-point rounding can put cosine slightly
 above 1. This single-image check establishes limited conversion parity, not retrieval
 accuracy across datasets. The supplied package retains its original conversion metadata.
 

SHA256SUMS

@@ -1,5 +1,5 @@
-bc41c5c4c50144420768edcc0651acfb7702fcfebf07722fd43d49ac8711cf00  models/DINOv3ViTB16CLS-FP32-448.mlpackage/Data/com.apple.CoreML/model.mlmodel
-12153d208093c1bfb72eed482d43a695e4062a7b654802355615145b088d5203  models/DINOv3ViTB16CLS-FP32-448.mlpackage/Data/com.apple.CoreML/weights/weight.bin
-3a4ef67c637d02b35ad1a16c1845c5a433b74eb7587be0a39f5ad751fd9fe483  models/DINOv3ViTB16CLS-FP32-448.mlpackage/Manifest.json
-be582eb7df4b60c1ef9df00e0134931edfa07a8d1506c2a35b1743084ff1f6ee  models/DINOv3ViTB16CLS-FP32-448.mlpackage/executorch_debug_handle_mapping.json
-97b94869da0008f660d3d2519a58602a4de51c5666b7ff44355d1c0395d447f2  models/DINOv3ViTB16CLS-FP32-448.validation.json
+a88e2a595ae587c134fa664280be243d2b9246643fa28e455601464125db7ea5  models/DINOv3ViTB16-FP32-448.mlpackage/Data/com.apple.CoreML/model.mlmodel
+12153d208093c1bfb72eed482d43a695e4062a7b654802355615145b088d5203  models/DINOv3ViTB16-FP32-448.mlpackage/Data/com.apple.CoreML/weights/weight.bin
+b080024a442c0a0d133905b7d79d25619451c784b10f456e0572dc81ac2270c9  models/DINOv3ViTB16-FP32-448.mlpackage/executorch_debug_handle_mapping.json
+a52ae1bc4b1203e88e0ff3355a347f8fbfbcad9150ea0df5073f16d1bfaca380  models/DINOv3ViTB16-FP32-448.mlpackage/Manifest.json
+9b01f646e006c36f84c5a8bdd4e7d006c805c25f0346380d97ff01506611a500  models/DINOv3ViTB16-FP32-448.validation.json

convert.py

@@ -17,18 +17,26 @@ IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD = (0.229, 0.224, 0.225)
 
 
-class DINOv3CLS(torch.nn.Module):
+class DINOv3Encoder(torch.nn.Module):
     def __init__(self, backbone: torch.nn.Module) -> None:
         super().__init__()
         self.backbone = backbone
+        self.num_prefix_tokens = backbone.num_prefix_tokens
         self.register_buffer("mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1))
         self.register_buffer("std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1))
 
-    def forward(self, image: torch.Tensor) -> torch.Tensor:
+    def forward(self, image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         pixels = (image / 255.0 - self.mean) / self.std
-        cls = self.backbone.forward_features(pixels)[:, 0]
+        features = self.backbone.forward_features(pixels)
+        cls = features[:, 0]
         norm = torch.sqrt(torch.sum(cls * cls, dim=-1, keepdim=True).clamp_min(1e-12))
-        return cls / norm
+        # The prefix is one CLS token plus four register tokens. Registers absorb
+        # high-norm artifacts that would otherwise pollute the patch tokens; they
+        # are not features and nothing downstream uses them.
+        patches = features[:, self.num_prefix_tokens:]
+        # Patches stay unnormalized because dense heads generally want the
+        # magnitude. Callers doing cosine can normalize per token themselves.
+        return cls / norm, patches
 
 
 def parse_args() -> argparse.Namespace:
@@ -57,7 +65,7 @@ def main() -> None:
         raise SystemExit("--size must be a positive multiple of 16")
     if args.output is None:
         precision_name = args.precision.upper()
-        args.output = Path(f"models/DINOv3ViTB16CLS-{precision_name}-{args.size}.mlpackage")
+        args.output = Path(f"models/DINOv3ViTB16-{precision_name}-{args.size}.mlpackage")
     if args.output.exists():
         if not args.force:
             raise SystemExit(f"{args.output} already exists; pass --force to replace it")
@@ -70,7 +78,7 @@ def main() -> None:
         img_size=args.size,
         num_classes=0,
     ).eval()
-    model = DINOv3CLS(backbone).eval()
+    model = DINOv3Encoder(backbone).eval()
 
     with torch.inference_mode():
         exported = torch.export.export(model, (example,)).run_decompositions({})
@@ -86,16 +94,17 @@ def main() -> None:
                 color_layout=ct.colorlayout.RGB,
             )
         ],
-        outputs=[ct.TensorType(name="embedding")],
+        outputs=[ct.TensorType(name="embedding"), ct.TensorType(name="patch_embeddings")],
         minimum_deployment_target=ct.target.macOS14,
         compute_precision=precision,
     )
     coreml_model.author = "dinov3-coreml; base model by Meta"
     coreml_model.license = "DINOv3 License"
-    coreml_model.short_description = "DINOv3 ViT-B/16 L2-normalized CLS embedding"
+    coreml_model.short_description = "DINOv3 ViT-B/16 CLS embedding and patch tokens"
     coreml_model.user_defined_metadata["base_model"] = BASE_MODEL
     coreml_model.input_description["image"] = f"RGB image resized to {args.size}x{args.size}"
     coreml_model.output_description["embedding"] = "L2-normalized 768-value CLS embedding"
+    coreml_model.output_description["patch_embeddings"] = "Unnormalized patch tokens, one per 16x16 patch"
 
     args.output.parent.mkdir(parents=True, exist_ok=True)
     coreml_model.save(args.output)
@@ -103,25 +112,32 @@ def main() -> None:
     image = synthetic_image(args.size)
     array = np.asarray(image, dtype=np.float32).transpose(2, 0, 1)[None, ...]
     with torch.inference_mode():
-        pytorch_output = model(torch.from_numpy(array)).numpy()[0]
-    coreml_output = np.asarray(coreml_model.predict({"image": image})["embedding"])[0]
-    cosine = float(
-        np.dot(pytorch_output, coreml_output)
-        / (np.linalg.norm(pytorch_output) * np.linalg.norm(coreml_output))
-    )
+        torch_cls, torch_patches = (t.numpy()[0] for t in model(torch.from_numpy(array)))
+    prediction = coreml_model.predict({"image": image})
+    coreml_output = np.asarray(prediction["embedding"])[0]
+    coreml_patches = np.asarray(prediction["patch_embeddings"])[0]
+
+    def cosine_of(a: np.ndarray, b: np.ndarray) -> float:
+        a, b = a.reshape(-1), b.reshape(-1)
+        return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
+
+    cosine = cosine_of(torch_cls, coreml_output)
+    patch_cosine = cosine_of(torch_patches, coreml_patches)
     report = {
         "model": DEFAULT_MODEL,
         "base_model": BASE_MODEL,
         "input_size": args.size,
         "precision": args.precision,
         "output_shape": list(coreml_output.shape),
+        "patch_output_shape": list(coreml_patches.shape),
         "pytorch_coreml_cosine_similarity": cosine,
+        "pytorch_coreml_patch_cosine_similarity": patch_cosine,
         "coreml_output_l2_norm": float(np.linalg.norm(coreml_output)),
     }
     report_path = args.output.with_suffix(".validation.json")
     report_path.write_text(json.dumps(report, indent=2) + "\n")
     print(json.dumps({"output": str(args.output), **report}, indent=2))
-    if cosine < 0.999:
+    if min(cosine, patch_cosine) < 0.999:
         raise SystemExit("Core ML parity check failed: cosine similarity is below 0.999")
 
 

examples/embed.py

@@ -11,7 +11,7 @@ from PIL import Image, ImageOps
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("image", type=Path, nargs="?")
-    parser.add_argument("--model", type=Path, default=Path("models/DINOv3ViTB16CLS-FP32-448.mlpackage"))
+    parser.add_argument("--model", type=Path, default=Path("models/DINOv3ViTB16-FP32-448.mlpackage"))
     parser.add_argument("--output", type=Path)
     args = parser.parse_args()
     model = ct.models.MLModel(str(args.model), compute_units=ct.ComputeUnit.ALL)
@@ -24,9 +24,13 @@ def main() -> None:
         y, x = np.mgrid[0:size[1], 0:size[0]]
         pixels = np.stack((x * 255 // size[0], y * 255 // size[1], (x + y) * 255 // sum(size)), axis=-1).astype(np.uint8)
         image = Image.fromarray(pixels)
-    vector = np.asarray(model.predict({"image": image})["embedding"], dtype=np.float32).reshape(-1)
+    prediction = model.predict({"image": image})
+    vector = np.asarray(prediction["embedding"], dtype=np.float32).reshape(-1)
+    patches = np.asarray(prediction["patch_embeddings"], dtype=np.float32).reshape(-1, 768)
     if vector.shape != (768,) or not np.isfinite(vector).all():
         raise SystemExit("Invalid embedding")
+    if not np.isfinite(patches).all():
+        raise SystemExit("Invalid patch embeddings")
     norm = float(np.linalg.norm(vector))
     if not np.isclose(norm, 1, atol=1e-4):
         raise SystemExit(f"Embedding is not normalized: {norm}")
@@ -34,6 +38,7 @@ def main() -> None:
         args.output.parent.mkdir(parents=True, exist_ok=True)
         np.save(args.output, vector)
     print(f"shape={vector.shape}, dtype={vector.dtype}, L2 norm={norm:.8f}")
+    print(f"patches={patches.shape} (unnormalized)")
 
 
 if __name__ == "__main__":

models/{DINOv3ViTB16CLS-FP32-448.mlpackage → DINOv3ViTB16-FP32-448.mlpackage}/Data/com.apple.CoreML/model.mlmodel

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bc41c5c4c50144420768edcc0651acfb7702fcfebf07722fd43d49ac8711cf00
-size 182130
+oid sha256:a88e2a595ae587c134fa664280be243d2b9246643fa28e455601464125db7ea5
+size 182835

models/{DINOv3ViTB16CLS-FP32-448.mlpackage → DINOv3ViTB16-FP32-448.mlpackage}/Manifest.json

@@ -1,18 +1,18 @@
 {
     "fileFormatVersion": "1.0.0",
     "itemInfoEntries": {
-        "0E300F4B-0BF4-477E-AE23-9C8EBE4EC7E8": {
-            "author": "com.apple.CoreML",
-            "description": "CoreML Model Weights",
-            "name": "weights",
-            "path": "com.apple.CoreML/weights"
-        },
-        "4C46342C-1CCD-429A-AAC6-E4BE6AD8B1D0": {
+        "5A08B9B3-DA0A-4DD2-9DBC-A68B509B51F6": {
             "author": "com.apple.CoreML",
             "description": "CoreML Model Specification",
             "name": "model.mlmodel",
             "path": "com.apple.CoreML/model.mlmodel"
+        },
+        "F9776A49-BC3F-4795-BFC2-9F9C12DEC863": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
         }
     },
-    "rootModelIdentifier": "4C46342C-1CCD-429A-AAC6-E4BE6AD8B1D0"
+    "rootModelIdentifier": "5A08B9B3-DA0A-4DD2-9DBC-A68B509B51F6"
 }

models/{DINOv3ViTB16CLS-FP32-448.validation.json → DINOv3ViTB16-FP32-448.validation.json}

@@ -6,6 +6,11 @@
   "output_shape": [
     768
   ],
+  "patch_output_shape": [
+    784,
+    768
+  ],
   "pytorch_coreml_cosine_similarity": 1.0000001192092896,
+  "pytorch_coreml_patch_cosine_similarity": 1.0000001192092896,
   "coreml_output_l2_norm": 0.9999999403953552
 }