\section{Methodology}

\begin{figure*}[t]
    \centering
    \includegraphics[width=\textwidth]{imgs/model_architecture10.jpg}
    \caption{Training framework for layout-aware ID representation learning. Top: self-supervised fine-tuning adapts the DINOv3 Vision Transformer to the document domain using SimMIM masked reconstruction. Bottom: Supervised metric learning maps ID images to 512-dimensional embeddings via a projection head, and the network is trained with a composite loss (ArcFace, supervised contrastive, and center loss) to encourage inter-class separability and intra-class compactness.}
    \label{fig:fig2}
\end{figure*}

We first adapt the general-purpose DINOv3 model to the document domain via context-aware self-supervised pre-training. Then, we further refine the learned representations via supervised metric learning with a composite loss function, yielding a highly discriminative embedding space suitable for downstream tasks such as classification, retrieval, clustering, and new set recognition. The training framework is shown in Figure \ref{fig:fig2}.

\subsection{Self-Supervised Fine-tuning}
Since the web-scale training of DINOv3 provides a robust, heavily generalizable prior for visual features, we adopt the pretrained DINOv3 weights as the initialization of our model's encoder. We use U.S. ID images for the training.
  
\paragraph{Image Augmentation} To improve robustness to real-world capture conditions and prevent overfitting, we apply strong data augmentation before the fine-tuning. The augmentation simulates imperfect framing, rotation, skew, and camera viewpoint changes commonly observed in mobile and handheld document capture. It preserves global layout semantics while introduces substantial appearance and geometric variability. We first resize each input image to 640x640. This is followed by mild color jittering, where brightness, contrast, and saturation are each independently perturbed within ±10\%, and hue is perturbed within ±0.02. The augmented images are then normalized by the channel-wise mean and standard deviation of ImageNet statistics \cite{deng2009imagenet}.

\paragraph{Model Architecture}  The DINOv3 encoder produces the patch embedding and the 1024-dimension global CLS token embedding for augmented ID images. We randomly mask 60\% of the patch embeddings, and use the remaining embeddings to reconstruct the original image. This design allows the self-attention mechanism to construct contextualized representations for all patches, informed by the global document layout. The reconstruction head is a Multi-Layer Perceptron (MLP) that predicts their original pixel values. 
  
\paragraph{Fine-tuning Strategy} To stabilize adaptation and reduce overfitting, we use progressive unfreezing. Initially, the encoder is frozen and only the reconstruction head is trained with AdamW. Starting at epoch 1, one transformer block per epoch (block 23 → 12) is unfrozen from top to bottom and added to the optimizer as a new parameter group. Training runs for 30 epochs, and the best model is selected by validation loss.

\subsection{Supervised Metric Learning}
We implement supervised learning to learn a structured embedding space that is highly discriminative for seen layouts, meanwhile generalizable to unseen layouts. We use 15k labeled IDs from almost all U.S. states. We split these layouts into disjoint training, validation, and test sets by states of origin. A batch size of 64 is used for training.

\paragraph{Data Augmentation} We first apply Geometric Augmentations, including random resized cropping (covering 60--100\% of the image area), affine transformations including rotation (up to $20^\circ$), translation (up to 15\%), shear, and probabilistic perspective distortion. We then apply Photometric Augmentations, including color jittering to brightness, contrast, saturation, and hue, followed by mild Gaussian blur to simulate variation in illumination, printing, and camera sensors. Finally, images are normalized using ImageNet statistics, and are randomly masked to discourage reliance on localized visual cues and encourage learning of global layout structure.

\paragraph{Model Architecture and Composite Loss Function}  We used the fine-tuned DINOv3 encoder (discard the reconstruction head) as the base. The encoder outputs a 1024-dimensional CLS token embedding, which is passed through a projection head implemented as a light-weighted MLP with batch normalization and dropout. The projection head reduces the dimensionality to a 512-dimensional embedding, which serves as the final representation for downstreaming tasks. To jointly encourage inter-class separability and intra-class compactness, we employ a composite loss function:
\begin{equation}
\mathcal{L}_{\text{total}} = \lambda_A \mathcal{L}_{\text{ArcFace}} + \lambda_S \mathcal{L}_{\text{SupCon}} + \lambda_C \mathcal{L}_{\text{Center}}
\end{equation}
where $\lambda_A$, $\lambda_S$, and $\lambda_C$ are the weights for the ArcFace loss, Supervised Contrastive loss, and Center loss, respectively. We will discuss the choice of these weights in the next section.
\paragraph{Training Strategy}  Our training strategy follows the widely adopted gradual unfreezing paradigm in transfer learning, where task-specific heads are first optimized while keeping the pretrained backbone frozen, followed by progressively unfreezing deeper layers. This staged approach has been shown to stabilize optimization and mitigate catastrophic forgetting when adapting pretrained models to new tasks \cite{howard2018universal, montone2017gradual}.
\begin{itemize}
\item Warm-up: The encoder backbone is frozen. Only the Projection and ArcFace heads are trained ($LR=5 \times 10^{-4}$). This establishes an initial cluster structure.
\item Partial Unfreezing: The last 8 blocks of the ViT backbone are unfrozen progressively ($LR=1 \times 10^{-5}$).
\item Full Fine-tuning: All layers are unfrozen with a reduced learning rate ($10\%$ of Partial Unfreezing) to refine the global feature extractors.
\end{itemize}
