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
|
from __future__ import annotations
import numpy as np
import ot
import torch
from sklearn.cluster import MiniBatchKMeans
from .common import pairwise_cosine_distance
def gw_pseudo_targets(
vision: torch.Tensor,
text: torch.Tensor,
clusters: int,
seed: int,
max_iter: int = 100,
) -> dict[str, torch.Tensor | float | int]:
"""Build structure-only pseudo-correspondences between modality prototypes."""
x = vision.float().numpy()
y = text.float().numpy()
k = min(clusters, len(x), len(y))
vx = MiniBatchKMeans(
n_clusters=k,
random_state=seed,
batch_size=min(4096, len(x)),
n_init=3,
max_iter=200,
).fit(x)
ty = MiniBatchKMeans(
n_clusters=k,
random_state=seed + 1,
batch_size=min(4096, len(y)),
n_init=3,
max_iter=200,
).fit(y)
cx = pairwise_cosine_distance(vx.cluster_centers_)
cy = pairwise_cosine_distance(ty.cluster_centers_)
p = np.bincount(vx.labels_, minlength=k).astype(np.float64)
q = np.bincount(ty.labels_, minlength=k).astype(np.float64)
p /= p.sum()
q /= q.sum()
coupling, log = ot.gromov.gromov_wasserstein(
cx,
cy,
p,
q,
loss_fun="square_loss",
armijo=False,
log=True,
max_iter=max_iter,
tol_rel=1e-8,
tol_abs=1e-8,
)
target_centers = coupling @ ty.cluster_centers_
target_centers /= p[:, None].clip(min=1e-12)
target_centers /= np.linalg.norm(target_centers, axis=1, keepdims=True).clip(
min=1e-12
)
row_entropy = -np.sum(
(coupling / p[:, None].clip(min=1e-12))
* np.log(
(coupling / p[:, None].clip(min=1e-12)).clip(min=1e-12)
),
axis=1,
).mean()
return {
"vision_centers": torch.from_numpy(vx.cluster_centers_).float(),
"text_centers": torch.from_numpy(ty.cluster_centers_).float(),
"target_centers": torch.from_numpy(target_centers).float(),
"vision_assignments": torch.from_numpy(vx.labels_).long(),
"coupling": torch.from_numpy(coupling).float(),
"gw_distance": float(log["gw_dist"]),
"coupling_row_entropy": float(row_entropy),
"clusters": k,
}
|
