Understanding Domain-Aware Distribution Alignment in Budgeted Entity Matching

Technical Breakdown§

Core Methodology The paper systematically studies BEACON, a framework for budgeted, domain-aware entity matching (EM). EM aims to identify records referring to the same real-world entity across different data sources (domains). BEACON leverages distribution alignment to transfer knowledge from a labeled source domain to a target domain with limited labels (budgeted setting). The core idea is to learn domain-invariant feature representations via adversarial training: a feature extractor $F$ is trained to fool a domain classifier $D$, while $D$ tries to predict which domain a sample comes from. This min-max game can be formalized as:

$$\mathcal{L}_{adv} = \mathbb{E}_{\mathbf{x} \sim p_{src}}[\log D(F(\mathbf{x}))] + \mathbb{E}_{\mathbf{x} \sim p_{tgt}}[\log (1 - D(F(\mathbf{x})))]$$

Simultaneously, a matching classifier $M$ is trained on source labels and a small budget of target labels:

$$\mathcal{L}_{cls} = \frac{1}{|\mathcal{D}_{src}|} \sum_{(\mathbf{x}_i, y_i) \in \mathcal{D}_{src}} \ell(M(F(\mathbf{x}_i)), y_i) + \frac{1}{|\mathcal{D}_{tgt}^{labeled}|} \sum_{(\mathbf{x}_j, y_j) \in \mathcal{D}_{tgt}^{labeled}} \ell(M(F(\mathbf{x}_j)), y_j)$$

The total loss is $\mathcal{L} = \mathcal{L}_{cls} + \lambda \mathcal{L}_{adv}$, where $\lambda$ is a trade-off parameter. The study varies data availability (number of target labels, source domain size) and algorithmic choices (backbone architecture, alignment strength) to understand their impact.

Implementation Details The framework is implemented as a PyTorch module with three components: a shared encoder (e.g., DistilBERT or BiLSTM), a domain classifier (a two-layer MLP with gradient reversal layer), and a matching classifier (MLP with softmax over three classes: match, non-match, uncertain). Training alternates between standard supervised learning and adversarial adaptation. The budgeted setting simulates active learning: initially only source domain labels are available; a small number of target examples are annotated iteratively.

Code Snippet

import torch
import torch.nn as nn
from torch.nn.utils import gradient_reversal

class FeatureExtractor(nn.Module):
    def __init__(self, input_dim, hidden_dim=256):
        super().__init__()
        self.encoder = nn.Linear(input_dim, hidden_dim)
        self.fc = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        return self.fc(torch.relu(self.encoder(x)))

class DomainClassifier(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.fc = nn.Linear(hidden_dim, 2)  # src vs tgt

    def forward(self, x, alpha=1.0):
        reversed_x = gradient_reversal(x, alpha)  # reverse gradient
        return torch.softmax(self.fc(reversed_x), dim=-1)

class MatchingClassifier(nn.Module):
    def __init__(self, hidden_dim, num_classes=3):
        super().__init__()
        self.fc = nn.Linear(hidden_dim, num_classes)

    def forward(self, x):
        return torch.softmax(self.fc(x), dim=-1)

# Training loop
for batch in dataloader:
    src_feat, tgt_feat = batch['src'], batch['tgt']
    features = torch.cat([F(src_feat), F(tgt_feat)], dim=0)
    domain_labels = torch.cat([torch.zeros(len(src_feat)), torch.ones(len(tgt_feat))], dim=0).long()
    
    # Adversarial loss
    domain_pred = D(features, alpha=0.1)
    loss_adv = nn.CrossEntropyLoss()(domain_pred, domain_labels)
    
    # Classification loss on labeled data
    src_pred = M(F(src_feat))
    loss_cls = nn.CrossEntropyLoss()(src_pred, src_labels)
    
    loss = loss_cls + lambda_ * loss_adv
    loss.backward()
    optimizer.step()

Key Findings

• Distribution alignment consistently improves EM performance when target labels are scarce (< 100).
• The benefit diminishes with larger target budgets, suggesting alignment is most critical under severe data constraints.
• Simple backbone architectures (e.g., BiLSTM) benefit more from alignment than large pretrained models (DistilBERT), likely because the latter already capture some domain-invariant features.
• The gradient reversal coefficient $\alpha$ must be tuned; too high can harm matching accuracy by losing domain-specific discriminative cues.