File size: 2,494 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | import torch
def alignment(
embeddings: torch.Tensor,
labels: torch.Tensor,
alpha: float = 2,
):
"""
https://arxiv.org/pdf/2005.10242
Label-aware Alignment loss.
Calculates alignment for embeddings of samples with the SAME label
within a batch, assuming embeddings are already unit-normalized.
Args:
embeddings: Tensor [N, D] - Batch of unit-normalized embeddings.
labels: Tensor [N] - Corresponding labels.
alpha: Power to raise squared distance (hyperparameter, default=2).
Returns:
Tensor: Label-aware Alignment loss (scalar). Returns 0 if no positive pairs.
"""
assert embeddings.size(0) == labels.size(0), "Embeddings and labels must have the same size."
n_samples = embeddings.size(0)
if n_samples < 2:
return torch.tensor(0.0, device=embeddings.device)
# Create a pairwise label comparison matrix (N x N), exclude self-pairs
labels_equal_mask = (labels[:, None] == labels[None, :]).triu(diagonal=1)
positive_indices = torch.nonzero(labels_equal_mask, as_tuple=False)
if positive_indices.numel() == 0:
return torch.tensor(0.0, device=embeddings.device)
# Get embeddings of positive pairs
x = embeddings[positive_indices[:, 0]]
y = embeddings[positive_indices[:, 1]]
# Calculate alignment loss
return (x - y).norm(p=2, dim=1).pow(alpha).mean()
def uniformity(
x: torch.Tensor,
t: float = 2,
clip_value: float = 1e-6,
):
"""
https://arxiv.org/pdf/2005.10242
Calculates the Uniformity loss.
Args:
x: [N, D] - Batch of feature embeddings.
t: Temperature parameter (hyperparameter).
Returns:
Tensor: Uniformity loss value (scalar).
"""
return torch.pdist(x, p=2).pow(2).mul(-t).exp().mean().clamp(min=clip_value).log()
if __name__ == "__main__":
embeddings = torch.tensor(
[
[1.0, 0.0],
[1.0, 0.0],
[1.0, 1.0],
[0.0, 1.0],
[0.0, 1.0],
],
)
embeddings /= embeddings.norm(p=2, dim=1, keepdim=True)
labels = torch.tensor([0, 0, 0, 1, 1])
print("Embeddings:")
print(embeddings.numpy())
print("\nLabels:")
print(labels.numpy())
alignment_loss = alignment(embeddings, labels, alpha=2)
print("\nAlignment loss:", alignment_loss.item())
uniformity_loss = uniformity(embeddings, t=2, clip_value=1e-6)
print("Uniformity loss:", uniformity_loss.item())
|