Re-thinking Mammography Transfer Learning: The Dataset-Informed Transfer Learning (DITL) Framework for Breast Cancer Screening and Lesion Diagnosis

Technical Synopsis§

The Dataset-Informed Transfer Learning (DITL) framework addresses the challenge of enhancing mammography classification by adapting transfer learning to dataset-specific characteristics. DITL integrates two novel components: Adaptive Difficulty-Weighted Cross-Entropy (A-DWCE) and Adaptive Neighborhood Representation Triplet (A-NR-Triplet).

Method Details§

Adaptive Difficulty-Weighted Cross-Entropy (A-DWCE) assigns per-sample weights based on the label purity of the sample's k-nearest neighbors in a self-supervised feature space. Given a sample $x_i$ with label $y_i$, let $\mathcal{N}_i$ be its k-nearest neighbor set. Define purity $p_i = \frac{1}{k} \sum_{j \in \mathcal{N}_i} \mathbb{I}(y_j = y_i)$. The weight $w_i$ is computed as $w_i = 1 - p_i$, so that samples with low purity (harder samples) receive higher weight. The loss becomes:

$$\mathcal{L}_{\text{A-DWCE}} = -\frac{1}{N} \sum_{i=1}^{N} w_i \log\left(\frac{e^{f_{y_i}(x_i)}}{\sum_{c} e^{f_c(x_i)}}\right)$$

where $f_c(x_i)$ is the logit for class $c$.

Adaptive Neighborhood Representation Triplet (A-NR-Triplet) enforces intra-class compactness and inter-class separation using a learnable margin. For each anchor $x_i$, a positive sample $x_p$ from the same class and a negative $x_n$ from a different class are selected from its neighborhood. The triplet loss with a learnable margin $m_i$ is:

$$\mathcal{L}_{\text{A-NR-Triplet}} = \max\left(0, d(f_i, f_p) - d(f_i, f_n) + m_i\right)$$

where $f_i$ are feature embeddings, and $m_i$ is adaptively set as $m_i = \alpha \cdot (1 - p_i)$ with $\alpha$ a scaling factor. The total loss is a weighted sum:

$$\mathcal{L}_{\text{DITL}} = \mathcal{L}_{\text{A-DWCE}} + \lambda \mathcal{L}_{\text{A-NR-Triplet}}$$

$\lambda$ is fixed (e.g., 1.0), and no hyperparameter tuning is required beyond standard training settings.

Implementation Details§

The framework is model-agnostic. A CNN backbone (e.g., ResNet-50) is pre-trained on ImageNet, then fine-tuned with DITL. Self-supervised features for computing nearest neighbors are obtained from a pre-trained SimCLR model on the target dataset. The neighborhood size $k$ is set to 10. The training uses standard SGD with momentum.

import torch
import torch.nn.functional as F

def ditl_loss(features, labels, logits, k=10, alpha=1.0, lambd=1.0):
    # features: [N, D] from backbone (without classifier)
    # logits: [N, C] from classifier head
    
    # Compute pairwise cosine similarity
    sim = F.normalize(features, dim=1) @ F.normalize(features, dim=1).T
    
    # Get k-nearest neighbors (exclude self)
    _, indices = torch.topk(sim, k=k+1, dim=1)
    neighbors = indices[:, 1:]  # remove self
    
    # Compute purity per sample
    labels_neighbors = labels[neighbors]  # [N, k]
    purity = (labels_neighbors == labels.unsqueeze(1)).float().mean(dim=1)  # [N]
    weights = 1 - purity
    
    # A-DWCE
    ce_loss = F.cross_entropy(logits, labels, reduction='none')
    weighted_ce = (weights * ce_loss).mean()
    
    # A-NR-Triplet
    # For each anchor, select positive: any neighbor with same label; negative: any neighbor with different label
    # Simplified: use first neighbor of same class as positive, first of different class as negative
    pos_mask = (labels[neighbors] == labels.unsqueeze(1)).float()  # [N, k]
    neg_mask = 1 - pos_mask
    # For simplicity, assume each anchor has both pos and neg neighbors; in practice handle edge cases
    pos_idx = (pos_mask.cumsum(dim=1) == 1).float().argmax(dim=1)  # first positive neighbor index
    neg_idx = (neg_mask.cumsum(dim=1) == 1).float().argmax(dim=1)  # first negative neighbor index
    
    feat = F.normalize(features, dim=1)
    anchor = feat
    positive = feat[neighbors.gather(1, pos_idx.unsqueeze(1)).squeeze()]
    negative = feat[neighbors.gather(1, neg_idx.unsqueeze(1)).squeeze()]
    
    d_pos = F.pairwise_distance(anchor, positive, p=2)
    d_neg = F.pairwise_distance(anchor, negative, p=2)
    margins = alpha * (1 - purity)
    triplet_loss = torch.clamp(d_pos - d_neg + margins, min=0).mean()
    
    total_loss = weighted_ce + lambd * triplet_loss
    return total_loss

Results§

On VinDR-Mammo (large-scale), DITL improves accuracy by 1.5%, F1 by 1.8%, AUC by 0.6% (p<0.0001) over standard fine-tuning. On small ROI datasets (CBIS-DDSM, INbreast), similar statistically significant gains are observed.