1. Prototype Attention Maps via Gumbel-Softmax
Given patch features \(\mathbf{z} \in \mathbb{R}^{D \times H \times W}\) from a pre-trained ViT, we compute attention maps \(\mathbf{A} \in [0, 1]^{(K+1) \times H \times W}\) representing \(K\) foreground parts and a background channel. Attention probabilities are computed by comparing patch tokens \(\mathbf{z}_{ij}\) with learnable prototypes \(\mathbf{p}^k\), sampled through a Gumbel-Softmax:
\[
a_{ij}^k = \frac{\exp\left(-\| \mathbf{z}_{ij} - \mathbf{p}^k\|^2 + \gamma_k \right)}{\sum_{l=1}^{K+1} \exp\left(-\| \mathbf{z}_{ij} - \mathbf{p}^l\|^2 + \gamma_l\right)}
\]
where \(\gamma_k \sim \text{Gumbel}(0, 1)\). Part embedding vectors \(\mathbf{v}^k\) are then obtained by weighted average pooling of features \(\mathbf{z}\) over each attention map \(A^k\).
2. Part Modulation & Dropout
Each part vector \(\mathbf{v}^k\) is modulated with a dedicated per-part Layer Normalization:
\[
\mathbf{v}_m^k = \frac{\mathbf{v}^k - \mathbb{E}[\mathbf{v}^k]}{\sqrt{\text{Var}[\mathbf{v}^k] + \epsilon}} \odot \mathbf{w}_m^k + \mathbf{b}_m^k
\]
During training, entire part embedding vectors are randomly dropped (Part Dropout). This forces each individual part to be discriminative for the fine-grained classification task rather than relying on a single dominant part.
3. Rethinking Geometric Part Priors: From Concentration to Total Variation
Previous weakly-supervised part discovery approaches (such as PDiscoNet, van der Klis et al.) enforce compactness on the learned attention maps \(\mathbf{A}^k\) through a concentration loss \(\mathcal{L}_{\text{conc}}\):
\[
\mathcal{L}_\text{conc} = \frac{\sum_{k=1}^K \sigma_v^2(\mathbf{A}^k) + \sigma_h^2(\mathbf{A}^k)}{K}
\]
where \(\sigma_v\) and \(\sigma_h\) denote the vertical and horizontal spatial variance respectively over the attention maps, computed with respect to the part centroid \((c_v^k, c_h^k)\):
\[
c_v^k = \frac{\sum_{i,j} i \cdot a_{ij}^k}{\sum_{i,j} a_{ij}^k}, \qquad
c_h^k = \frac{\sum_{i,j} j \cdot a_{ij}^k}{\sum_{i,j} a_{ij}^k}
\]
\[
\sigma_v^2(\mathbf{A}^k) = \sum_{i,j} a_{ij}^k \left( \frac{i - c_v^k}{H} \right)^2, \qquad
\sigma_h^2(\mathbf{A}^k) = \sum_{i,j} a_{ij}^k \left( \frac{j - c_h^k}{W} \right)^2
\]
While \(\mathcal{L}_{\text{conc}}\) was essential for CNN backbones to avoid degenerate solutions, penalizing spatial variance strictly constrains discovered parts to be compact and unimodal. As a result, it fundamentally struggles with:
- Multi-component parts: Natural parts that are spatially separated (e.g., both wings of a bird, pairs of quadruped legs, disconnected flower petals).
- Irregular and elongated geometries: Non-compact structures such as snake bodies, stems, or animal tails.
Because self-supervised Vision Transformers (such as DINOv2) already encode strong patch-level semantic grouping, PDiscoFormer relaxes this rigid centroid constraint by replacing \(\mathcal{L}_{\text{conc}}\) with a Total Variation prior:
-
Total Variation Loss (\(\mathcal{L}_{\text{tv}}\)):
Penalizes spatial gradients of attention maps to favor piece-wise constant, spatially connected regions without restricting shape, elongation, or number of connected components:
\[
\mathcal{L}_{\text{tv}} = \frac{1}{HW} \sum_{k=1}^{K+1} \sum_{ij} \left| \nabla a_{ij}^k \right|
\]
-
Entropy Loss (\(\mathcal{L}_{\text{ent}}\)):
Encourages crisp assignment of each patch to a single part prototype:
\[
\mathcal{L}_{\text{ent}} = \frac{-1}{K+1}\sum_{k=1}^{K+1}\sum_{ij} a_{ij}^k \log\left(a_{ij}^k\right)
\]
-
Background Presence Loss (\(\mathcal{L}_{\text{p}_0}\)):
Enforces the background channel to be present in every image, prioritized near image boundaries via a soft spatial border mask \(m_{ij}\):
\[
\mathcal{L}_{\text{p}_0} = -\frac{1}{B}\sum_b \log\left(\max_{i,j} m_{ij} \bar{a}_{ij}^{K+1}(\mathbf{x}_b)\right)
\]
-
Orthogonality (\(\mathcal{L}_{\perp}\)) & Equivariance (\(\mathcal{L}_{\text{eq}}\)):
Ensure that different part embedding vectors are decorrelated, and that predicted part maps transform equivariantly under rigid affine augmentations.
