Add Mettle RC2 weights, loader, and model card

README.md

@@ -0,0 +1,145 @@
+---
+license: other
+base_model: bioptimus/H-optimus-0
+library_name: transformers
+pipeline_tag: feature-extraction
+tags:
+  - computational-pathology
+  - histopathology
+  - foundation-model
+  - feature-extraction
+  - scanner-robustness
+---
+
+# Mettle
+
+Mettle is a histopathology tile encoder derived from H-optimus-0. It is designed to improve representation consistency across slide scanners and staining variation while retaining biologically useful morphology.
+
+Mettle is intended for research feature extraction. It is not a diagnostic system and has not been validated for autonomous clinical use.
+
+## Model details
+
+| Item | Value |
+|---|---|
+| Parent model | [H-optimus-0](https://huggingface.co/bioptimus/H-optimus-0) |
+| Parent-model license | Apache-2.0 |
+| Architecture | ViT-g/14, DINOv2 reg4 |
+| Input | 224 × 224 RGB histology tile |
+| Compatibility output | 1,536-dimensional headed CLS embedding |
+| Public benchmark output | 3,072-dimensional CLS+mean-patch embedding |
+| `model.safetensors` SHA-256 | `7754ef463e79dd13b072247b56db39791c884db10774e354d76f76367093ce4a` |
+
+The public benchmark representation concatenates the headed CLS embedding with the mean of the spatial patch tokens:
+
+```text
+[headed CLS embedding (1536), mean spatial-patch embedding (1536)]
+```
+
+CLS and register tokens are excluded from the patch mean. The first 1,536 coordinates are exactly the compatibility CLS representation.
+
+## Public benchmark results
+
+All results in this table use the 3,072-dimensional CLS+mean representation. Higher is better. Values after `±` are standard deviations over five downstream probe runs. PathoROB Robustness Index is a point estimate without an uncertainty interval.
+
+| Benchmark | Metric | Mettle | H-optimus-0 |
+|---|---|---:|---:|
+| PathoROB | Mean Robustness Index | **0.92227** | 0.8015 |
+| HEST | Mean Pearson, 9 tasks | **0.43032** | 0.42526 |
+| EVA BACH | Accuracy | **0.78168** ± 0.00567 | 0.7294 |
+| EVA CRC | Accuracy | **0.96912** | 0.9578 |
+| EVA PatchCamelyon | Balanced accuracy | 0.9476 | 0.9441 |
+
+### PathoROB breakdown
+
+| Cohort | Mettle |
+|---|---:|
+| Camelyon | 0.93859 |
+| TCGA | 0.86521 |
+| Tolkach ESCA | 0.96302 |
+| **Mean** | **0.92227** |
+
+PathoROB does not provide uncertainty estimates for its Robustness Index. Differences in point estimates should not be described as statistical superiority.
+
+### HEST task breakdown
+
+HEST uses PCA to 256 dimensions followed by ridge regression for a 50-gene panel. The metric is Pearson correlation averaged over official splits.
+
+| HEST task | Mettle | Published H-optimus-0 |
+|---|---:|---:|
+| IDC | 0.6104 | **0.6106** |
+| PRAD | **0.3789** | 0.36215 |
+| PAAD | **0.5309** | 0.5106 |
+| SKCM | **0.6745** | 0.6614 |
+| COAD | **0.3331** | 0.3089 |
+| READ | 0.2372 | **0.2401** |
+| CCRCC | 0.2552 | **0.2919** |
+| LUNG | **0.5777** | 0.5754 |
+| LYMPH_IDC | **0.2751** | 0.2665 |
+| **Unweighted mean** | **0.4303** | **0.4253** |
+
+## Usage
+
+Install the dependencies listed in `requirements.txt`, then load the image processor and model through the standard Transformers auto classes.
+
+```python
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModel
+
+model_id = "slideflow-labs/Mettle"
+processor = AutoImageProcessor.from_pretrained(model_id)
+model = AutoModel.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+).eval().cuda()
+
+image = Image.open("tile.png").convert("RGB")
+pixel_values = processor(
+    images=image,
+    return_tensors="pt",
+)["pixel_values"].cuda()
+
+with torch.inference_mode():
+    # Representation used for the public benchmarks: (batch, 3072)
+    cls_mean = model.encode(pixel_values, feature_view="cls_mean")
+
+    # Conventional compatibility representation: (batch, 1536)
+    cls = model.encode(pixel_values, feature_view="cls")
+```
+
+## Choosing a representation
+
+The reported public benchmarks use the 3,072-dimensional CLS+mean representation. The loader also exposes the 1,536-dimensional CLS embedding for compatibility; CLS can also be stronger on some endpoints.
+
+## Intended use
+
+Mettle is intended for research involving histopathology representation learning, feature extraction, model development, and robustness analysis. Users should validate the model on the intended tissue, acquisition hardware, population, and downstream endpoint.
+
+Do not use Mettle as a standalone diagnostic system. Appropriate human oversight, validation, and compliance requirements remain necessary.
+
+## Limitations
+
+- Performance can vary with tissue type, preparation, stain, scanner, compression, magnification, and downstream protocol.
+- Scanner consistency does not guarantee preservation of every biological signal or subgroup.
+- The two output representations are endpoint-dependent; neither dominates every benchmark.
+- PathoROB Robustness Index is reported without uncertainty estimates.
+- Benchmark results do not establish clinical efficacy or safety.
+- The parent model's limitations and license obligations continue to apply.
+
+## References
+
+- [H-optimus-0 model card](https://huggingface.co/bioptimus/H-optimus-0)
+- [Towards robust foundation models for digital pathology
+  (PathoROB)](https://doi.org/10.1038/s41467-026-73923-2)
+- [HEST-1k: A Dataset for Spatial Transcriptomics and Histology Image
+  Analysis](https://papers.neurips.cc/paper_files/paper/2024/file/60a899cc31f763be0bde781a75e04458-Paper-Datasets_and_Benchmarks_Track.pdf)
+- [eva: Evaluation framework for pathology foundation
+  models](https://openreview.net/forum?id=FNBQOPj18N)
+
+## Citation
+
+An arXiv preprint is forthcoming.
+
+## License
+
+Mettle is licensed under CC-BY-NC-ND-4.0, and is available for academic, non-commercial use.

USAGE.txt

@@ -0,0 +1,28 @@
+Mettle usage
+============
+
+Install the dependencies in requirements.txt, then load a local checkout or
+Hugging Face repository with the standard Transformers auto classes:
+
+    from PIL import Image
+    from transformers import AutoImageProcessor, AutoModel
+
+    repo = "slideflow-labs/Mettle"
+    processor = AutoImageProcessor.from_pretrained(repo)
+    model = AutoModel.from_pretrained(
+        repo,
+        trust_remote_code=True,
+    ).eval()
+
+    image = Image.open("tile.png").convert("RGB")
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+
+    # Public benchmark representation: (batch, 3072)
+    embedding = model.encode(pixel_values, feature_view="cls_mean")
+
+    # Compatibility representation: (batch, 1536)
+    cls_embedding = model.encode(pixel_values, feature_view="cls")
+
+The first 1536 coordinates of cls_mean are exactly cls_embedding. The remaining
+1536 coordinates are the mean of spatial patch tokens; CLS and register tokens
+are excluded.

__init__.py

@@ -0,0 +1,6 @@
+"""Mettle Hugging Face custom model."""
+
+from .configuration_mettle import MettleConfig
+from .modeling_mettle import MettleModel, MettleModelOutput
+
+__all__ = ["MettleConfig", "MettleModel", "MettleModelOutput"]

config.json

@@ -0,0 +1,20 @@
+{
+  "architectures": [
+    "MettleModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mettle.MettleConfig",
+    "AutoModel": "modeling_mettle.MettleModel"
+  },
+  "backbone_name": "vit_giant_patch14_reg4_dinov2",
+  "default_feature_view": "cls",
+  "head_enabled": true,
+  "head_hidden_size": 256,
+  "head_num_atoms": 16,
+  "head_rank": 8,
+  "hidden_size": 1536,
+  "image_size": 224,
+  "model_type": "mettle",
+  "num_register_tokens": 4,
+  "torch_dtype": "float32"
+}

configuration_mettle.py

@@ -0,0 +1,38 @@
+"""Configuration for the Mettle pathology tile encoder."""
+
+from transformers import PretrainedConfig
+
+
+class MettleConfig(PretrainedConfig):
+    """Configuration for a Mettle ViT-g/14 encoder."""
+
+    model_type = "mettle"
+
+    def __init__(
+        self,
+        backbone_name="vit_giant_patch14_reg4_dinov2",
+        image_size=224,
+        hidden_size=1536,
+        num_register_tokens=4,
+        head_enabled=True,
+        head_num_atoms=16,
+        head_rank=8,
+        head_hidden_size=256,
+        default_feature_view="cls",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if default_feature_view not in {"cls", "cls_mean"}:
+            raise ValueError(
+                "default_feature_view must be 'cls' or 'cls_mean', got "
+                f"{default_feature_view!r}"
+            )
+        self.backbone_name = backbone_name
+        self.image_size = int(image_size)
+        self.hidden_size = int(hidden_size)
+        self.num_register_tokens = int(num_register_tokens)
+        self.head_enabled = bool(head_enabled)
+        self.head_num_atoms = int(head_num_atoms)
+        self.head_rank = int(head_rank)
+        self.head_hidden_size = int(head_hidden_size)
+        self.default_feature_view = default_feature_view

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7754ef463e79dd13b072247b56db39791c884db10774e354d76f76367093ce4a
+size 4545473592

modeling_mettle.py

@@ -0,0 +1,163 @@
+"""Hugging Face model implementation for the Mettle tile encoder."""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import timm
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+try:
+    from .configuration_mettle import MettleConfig
+except ImportError:
+    from configuration_mettle import MettleConfig
+
+
+class MettleRefineHead(nn.Module):
+    """Content-conditional refinement of the pooled CLS embedding."""
+
+    def __init__(self, dim, num_atoms, rank, hidden_size):
+        super().__init__()
+        self.P = nn.Parameter(torch.zeros(num_atoms, dim, rank))
+        self.Q = nn.Parameter(torch.zeros(num_atoms, dim, rank))
+        self.mix = nn.Sequential(
+            nn.Linear(dim, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, num_atoms),
+        )
+        self.shift = nn.Sequential(
+            nn.Linear(dim, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, dim),
+        )
+
+    def forward(self, embedding):
+        mixture = self.mix(embedding)
+        shift = self.shift(embedding)
+        projected = torch.einsum("bd,kdr->bkr", embedding, self.P)
+        delta = torch.einsum("bkr,kdr->bkd", projected, self.Q)
+        return embedding + (mixture.unsqueeze(-1) * delta).sum(1) + shift
+
+
+@dataclass
+class MettleModelOutput(ModelOutput):
+    """Outputs from :class:`MettleModel`.
+
+    ``pooler_output`` follows the requested feature view. ``cls_embedding`` is
+    the headed 1536-dimensional CLS representation. ``cls_mean_embedding`` is
+    the concatenation of headed CLS and the unheaded mean spatial patch token.
+    """
+
+    last_hidden_state: Optional[torch.Tensor] = None
+    pooler_output: Optional[torch.Tensor] = None
+    cls_embedding: Optional[torch.Tensor] = None
+    mean_patch_embedding: Optional[torch.Tensor] = None
+    cls_mean_embedding: Optional[torch.Tensor] = None
+
+
+class MettleModel(PreTrainedModel):
+    """Mettle scanner-resilient pathology tile encoder."""
+
+    config_class = MettleConfig
+    base_model_prefix = "backbone"
+    main_input_name = "pixel_values"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone = timm.create_model(
+            config.backbone_name,
+            pretrained=False,
+            num_classes=0,
+            init_values=1e-5,
+            dynamic_img_size=False,
+            img_size=config.image_size,
+        )
+        if int(self.backbone.num_features) != config.hidden_size:
+            raise ValueError(
+                f"{config.backbone_name} has width {self.backbone.num_features}, "
+                f"but the config declares {config.hidden_size}"
+            )
+        self.head = None
+        if config.head_enabled:
+            self.head = MettleRefineHead(
+                dim=config.hidden_size,
+                num_atoms=config.head_num_atoms,
+                rank=config.head_rank,
+                hidden_size=config.head_hidden_size,
+            )
+        self.post_init()
+
+    def _embeddings(self, pixel_values):
+        tokens = self.backbone.forward_features(pixel_values)
+        expected_prefix = 1 + self.config.num_register_tokens
+        actual_prefix = int(self.backbone.num_prefix_tokens)
+        if actual_prefix != expected_prefix:
+            raise RuntimeError(
+                f"expected {expected_prefix} prefix tokens, found {actual_prefix}"
+            )
+        patches = tokens[:, actual_prefix:]
+        if patches.shape[1] == 0:
+            raise RuntimeError("the backbone returned no spatial patch tokens")
+
+        cls_embedding = tokens[:, 0].to(torch.float32)
+        if self.head is not None:
+            cls_embedding = self.head(cls_embedding)
+        mean_patch_embedding = patches.to(torch.float32).mean(dim=1)
+        cls_mean_embedding = torch.cat(
+            (cls_embedding, mean_patch_embedding),
+            dim=-1,
+        )
+        return tokens, cls_embedding, mean_patch_embedding, cls_mean_embedding
+
+    def forward(
+        self,
+        pixel_values,
+        feature_view=None,
+        return_dict=None,
+    ):
+        """Encode normalized 224-pixel RGB tiles.
+
+        Args:
+            pixel_values: Tensor with shape ``(batch, 3, 224, 224)`` after the
+                preprocessing specified by ``preprocessor_config.json``.
+            feature_view: ``"cls"`` for a 1536-dimensional embedding or
+                ``"cls_mean"`` for headed CLS concatenated with mean spatial
+                patch pooling (3072 dimensions).
+            return_dict: Follow the standard Transformers return convention.
+        """
+        feature_view = feature_view or self.config.default_feature_view
+        if feature_view not in {"cls", "cls_mean"}:
+            raise ValueError(
+                "feature_view must be 'cls' or 'cls_mean', got "
+                f"{feature_view!r}"
+            )
+        return_dict = (
+            self.config.use_return_dict if return_dict is None else return_dict
+        )
+        tokens, cls_embedding, mean_patch, cls_mean = self._embeddings(
+            pixel_values
+        )
+        selected = cls_embedding if feature_view == "cls" else cls_mean
+        if not return_dict:
+            return tokens, selected, cls_embedding, mean_patch, cls_mean
+        return MettleModelOutput(
+            last_hidden_state=tokens,
+            pooler_output=selected,
+            cls_embedding=cls_embedding,
+            mean_patch_embedding=mean_patch,
+            cls_mean_embedding=cls_mean,
+        )
+
+    def encode(self, pixel_values, feature_view=None):
+        """Return only the selected tile embedding."""
+        return self(
+            pixel_values=pixel_values,
+            feature_view=feature_view,
+            return_dict=True,
+        ).pooler_output
+
+    def get_input_embeddings(self):
+        """Return the ViT patch projection for Transformers integrations."""
+        return self.backbone.patch_embed.proj

package_model.py

@@ -0,0 +1,141 @@
+"""Convert a flat Mettle checkpoint into a Hugging Face safetensors artifact."""
+
+import argparse
+import hashlib
+import json
+import os
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+
+
+EXPECTED_HEAD_KEYS = {
+    "head.P",
+    "head.Q",
+    "head.mix.0.bias",
+    "head.mix.0.weight",
+    "head.mix.2.bias",
+    "head.mix.2.weight",
+    "head.shift.0.bias",
+    "head.shift.0.weight",
+    "head.shift.2.bias",
+    "head.shift.2.weight",
+}
+RC2_SOURCE_SHA256 = (
+    "ad79736b61638a58a72c5dedd7eb21582e94c6397564ae28c1c5da0af4d8a9be"
+)
+
+
+def sha256(path, chunk_size=16 * 1024 * 1024):
+    digest = hashlib.sha256()
+    with open(path, "rb") as stream:
+        while chunk := stream.read(chunk_size):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def convert_state_dict(source):
+    """Prefix backbone tensors for :class:`MettleModel`."""
+    head_keys = {key for key in source if key.startswith("head.")}
+    if head_keys != EXPECTED_HEAD_KEYS:
+        missing = sorted(EXPECTED_HEAD_KEYS - head_keys)
+        unexpected = sorted(head_keys - EXPECTED_HEAD_KEYS)
+        raise ValueError(
+            "checkpoint refinement-head keys do not match the release contract; "
+            f"missing={missing}, unexpected={unexpected}"
+        )
+    converted = OrderedDict()
+    for key in sorted(source):
+        target = key if key.startswith("head.") else f"backbone.{key}"
+        converted[target] = source[key].contiguous()
+    return converted
+
+
+def validate_shapes(state):
+    p_shape = tuple(state["head.P"].shape)
+    q_shape = tuple(state["head.Q"].shape)
+    if p_shape != (16, 1536, 8) or q_shape != p_shape:
+        raise ValueError(
+            "checkpoint refinement-head shape does not match RC2: "
+            f"P={p_shape}, Q={q_shape}"
+        )
+    if tuple(state["head.mix.0.weight"].shape) != (256, 1536):
+        raise ValueError("checkpoint refinement-head hidden size is not 256")
+    if tuple(state["head.shift.2.weight"].shape) != (1536, 256):
+        raise ValueError("checkpoint refinement-head output width is not 1536")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint", help="flat RC2 .pt checkpoint")
+    parser.add_argument(
+        "--output-dir",
+        default=str(Path(__file__).resolve().parent),
+        help="directory containing the Hugging Face repository files",
+    )
+    parser.add_argument(
+        "--expected-sha256",
+        default=RC2_SOURCE_SHA256,
+        help="refuse a source checkpoint whose SHA-256 differs",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="validate and map the checkpoint without writing weights",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    checkpoint = Path(args.checkpoint).resolve()
+    output_dir = Path(args.output_dir).resolve()
+    source_sha = sha256(checkpoint)
+    if args.expected_sha256 and source_sha != args.expected_sha256.lower():
+        raise SystemExit(
+            f"source SHA-256 mismatch: expected {args.expected_sha256}, "
+            f"found {source_sha}"
+        )
+
+    source = torch.load(checkpoint, map_location="cpu", weights_only=True)
+    validate_shapes(source)
+    converted = convert_state_dict(source)
+    print(
+        f"validated {len(source)} tensors; mapped {len(converted)} tensors; "
+        f"source sha256={source_sha}"
+    )
+    if args.dry_run:
+        return
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    output_path = output_dir / "model.safetensors"
+    manifest_path = output_dir / "weights_manifest.json"
+    if output_path.exists():
+        raise SystemExit(f"refusing to overwrite {output_path}")
+    if manifest_path.exists():
+        raise SystemExit(f"refusing to overwrite {manifest_path}")
+    temporary = output_dir / ".model.safetensors.incomplete"
+    if temporary.exists():
+        raise SystemExit(f"refusing to overwrite incomplete file {temporary}")
+    save_file(converted, str(temporary), metadata={"format": "pt"})
+    os.replace(temporary, output_path)
+    packaged_sha = sha256(output_path)
+    manifest = {
+        "format": "safetensors",
+        "model_file": output_path.name,
+        "model_sha256": packaged_sha,
+        "source_checkpoint_sha256": source_sha,
+        "tensor_count": len(converted),
+    }
+    manifest_path.write_text(
+        json.dumps(manifest, indent=2, sort_keys=True) + "\n",
+        encoding="utf-8",
+    )
+    print(f"wrote {output_path} sha256={packaged_sha}")
+    print(f"wrote {manifest_path}")
+
+
+if __name__ == "__main__":
+    main()

preprocessor_config.json

@@ -0,0 +1,23 @@
+{
+  "do_center_crop": false,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.707223,
+    0.578729,
+    0.703617
+  ],
+  "image_processor_type": "ViTImageProcessor",
+  "image_std": [
+    0.211883,
+    0.230117,
+    0.177517
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "height": 224,
+    "width": 224
+  }
+}

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.2
+timm>=1.0.15
+transformers>=4.48
+safetensors>=0.4

run_test.py

@@ -0,0 +1,44 @@
+"""Smoke-test a packaged local or downloaded Mettle-MX repository."""
+
+import argparse
+
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModel
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("model", help="local directory or Hugging Face repository ID")
+    parser.add_argument("--image", help="optional RGB image to encode")
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    processor = AutoImageProcessor.from_pretrained(args.model)
+    model = AutoModel.from_pretrained(
+        args.model,
+        trust_remote_code=True,
+    ).to(args.device).eval()
+    if args.image:
+        image = Image.open(args.image).convert("RGB")
+    else:
+        image = Image.new("RGB", (224, 224), color=(128, 96, 128))
+    inputs = processor(images=image, return_tensors="pt")
+    pixel_values = inputs["pixel_values"].to(args.device)
+    with torch.inference_mode():
+        cls = model.encode(pixel_values, feature_view="cls")
+        cls_mean = model.encode(pixel_values, feature_view="cls_mean")
+    if tuple(cls.shape) != (1, 1536):
+        raise RuntimeError(f"unexpected CLS shape: {tuple(cls.shape)}")
+    if tuple(cls_mean.shape) != (1, 3072):
+        raise RuntimeError(f"unexpected CLS+mean shape: {tuple(cls_mean.shape)}")
+    if not torch.equal(cls, cls_mean[:, :1536]):
+        raise RuntimeError("CLS differs between the two feature views")
+    print(f"PASS cls={tuple(cls.shape)} cls_mean={tuple(cls_mean.shape)}")
+
+
+if __name__ == "__main__":
+    main()

verify_against_flat.py

@@ -0,0 +1,119 @@
+"""Numerically compare the release model with a flat Mettle checkpoint."""
+
+import argparse
+import gc
+from pathlib import Path
+
+import timm
+import torch
+
+from configuration_mettle import MettleConfig
+from modeling_mettle import MettleModel, MettleRefineHead
+from package_model import convert_state_dict, validate_shapes
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint", help="flat Mettle checkpoint")
+    parser.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+    )
+    parser.add_argument("--atol", type=float, default=1e-5)
+    return parser.parse_args()
+
+
+def load_config():
+    return MettleConfig.from_json_file(
+        str(Path(__file__).with_name("config.json"))
+    )
+
+
+def reference_outputs(config, state, pixel_values):
+    backbone = timm.create_model(
+        config.backbone_name,
+        pretrained=False,
+        num_classes=0,
+        init_values=1e-5,
+        dynamic_img_size=False,
+        img_size=config.image_size,
+    )
+    backbone_state = {
+        key: value for key, value in state.items()
+        if not key.startswith("head.")
+    }
+    backbone.load_state_dict(backbone_state, strict=True)
+    head = MettleRefineHead(
+        dim=config.hidden_size,
+        num_atoms=config.head_num_atoms,
+        rank=config.head_rank,
+        hidden_size=config.head_hidden_size,
+    )
+    head_state = {
+        key.removeprefix("head."): value for key, value in state.items()
+        if key.startswith("head.")
+    }
+    head.load_state_dict(head_state, strict=True)
+    backbone.to(pixel_values.device).eval()
+    head.to(pixel_values.device).eval()
+    with torch.inference_mode():
+        tokens = backbone.forward_features(pixel_values)
+        cls = head(tokens[:, 0].to(torch.float32))
+        mean_patch = tokens[:, int(backbone.num_prefix_tokens):].float().mean(1)
+        cls_mean = torch.cat((cls, mean_patch), dim=-1)
+    return cls.cpu(), cls_mean.cpu()
+
+
+def release_outputs(config, state, pixel_values):
+    model = MettleModel(config)
+    model.load_state_dict(convert_state_dict(state), strict=True)
+    model.to(pixel_values.device).eval()
+    with torch.inference_mode():
+        cls = model.encode(pixel_values, "cls")
+        cls_mean = model.encode(pixel_values, "cls_mean")
+    return cls.cpu(), cls_mean.cpu()
+
+
+def main():
+    args = parse_args()
+    config = load_config()
+    state = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
+    validate_shapes(state)
+    generator = torch.Generator().manual_seed(20260730)
+    pixel_values = torch.randn(
+        1,
+        3,
+        config.image_size,
+        config.image_size,
+        generator=generator,
+    ).to(args.device)
+
+    reference_cls, reference_cls_mean = reference_outputs(
+        config,
+        state,
+        pixel_values,
+    )
+    if args.device.startswith("cuda"):
+        torch.cuda.empty_cache()
+    gc.collect()
+    release_cls, release_cls_mean = release_outputs(
+        config,
+        state,
+        pixel_values,
+    )
+
+    cls_delta = (reference_cls - release_cls).abs().max().item()
+    cls_mean_delta = (
+        reference_cls_mean - release_cls_mean
+    ).abs().max().item()
+    print(f"CLS max absolute delta: {cls_delta:.3e}")
+    print(f"CLS+mean max absolute delta: {cls_mean_delta:.3e}")
+    if cls_delta > args.atol or cls_mean_delta > args.atol:
+        raise SystemExit(
+            "FAILED: release outputs differ from the flat-checkpoint reference"
+        )
+    print("PASS: both feature views reproduce the flat-checkpoint reference")
+
+
+if __name__ == "__main__":
+    main()