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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
|
"""A matching benchmark, and the solvers we never tried on it.
Everything the project has used to search is one narrow family: spectral
initialisation on the pairwise field, exact steepest descent over
transpositions, tempering, and Sinkhorn. All of them act on the permutation
group with a pairwise objective in flat space. The measured cost of that
narrowness is large and specific -- on a field whose anchor bound is 0.997,
meaning the information is intact almost perfectly, blind matching returns
5.6%.
That makes a clean benchmark, because unlike natural data the answer is known
to be reachable. Three synthetic fields span the range: one that already
recovers, one whose captions omit a factor, one truncated to rank eight.
Natural data is included as the case where no method should be expected to
work, so a method that scores there is reporting a bug rather than a result.
Gromov-Wasserstein is the canonical formulation for exactly our problem --
align two metric spaces with no shared embedding, by matching their internal
distance structures -- and it has never been run on these fields. Its entropic
form with an annealed regulariser is the landscape-deformation method that the
rank ladder argued for and that we never reached. Each solver is scored blind
and then re-scored after exact refinement, since composition is what made the
spectral solver work.
"""
from __future__ import annotations
import argparse
import json
import time
import numpy as np
import torch
from .common import write_json
from .spectral_match import grampa, umeyama
from .synth_fast_gate import ClosedFormEnergy, fast_pair_descent
from .synth_triangle_gate import standardized
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--fields", nargs="+", required=True)
parser.add_argument("--labels", nargs="+", default=None)
parser.add_argument("--trials", type=int, default=3)
parser.add_argument("--iterations", type=int, default=600)
parser.add_argument("--device", default="cuda:3")
parser.add_argument("--methods", nargs="+", default=None)
parser.add_argument("--output", default="artifacts/synth_v1/match_battery.json")
return parser.parse_args()
def standardise(matrix: np.ndarray) -> np.ndarray:
mask = ~np.eye(len(matrix), dtype=bool)
values = matrix[mask]
out = (matrix - values.mean()) / values.std()
np.fill_diagonal(out, 0.0)
return out
def coupling_to_permutation(coupling: np.ndarray) -> np.ndarray:
from scipy.optimize import linear_sum_assignment
_, columns = linear_sum_assignment(-coupling)
return columns
# --- solvers: each returns column assignment for rows of `visual` ---------
def solve_grampa(visual, shuffled, **_):
return grampa(visual, shuffled, 1.0)
def solve_umeyama(visual, shuffled, **_):
return umeyama(visual, shuffled)
def solve_gw(visual, shuffled, **_):
import ot
size = len(visual)
weights = ot.unif(size)
coupling = ot.gromov.gromov_wasserstein(
visual, shuffled, weights, weights, "square_loss",
max_iter=200, tol_rel=1e-9,
)
return coupling_to_permutation(coupling)
def solve_gw_annealed(visual, shuffled, **_):
"""Entropic GW with the regulariser annealed from smooth to sharp."""
import ot
size = len(visual)
weights = ot.unif(size)
coupling = None
for epsilon in (5e-2, 2e-2, 1e-2, 5e-3, 2e-3):
try:
coupling = ot.gromov.entropic_gromov_wasserstein(
visual, shuffled, weights, weights, "square_loss",
epsilon=epsilon, G0=coupling, max_iter=200, tol=1e-9,
solver="PGD",
)
except Exception:
break
if not np.isfinite(coupling).all():
break
if coupling is None or not np.isfinite(coupling).all():
raise RuntimeError("entropic GW diverged at every regulariser")
return coupling_to_permutation(coupling)
def solve_bapg(visual, shuffled, **_):
"""Bregman alternating projected gradient: a newer GW solver."""
import ot
size = len(visual)
weights = ot.unif(size)
coupling = ot.gromov.BAPG_gromov_wasserstein(
visual, shuffled, weights, weights, "square_loss",
epsilon=1e-1, max_iter=1000, tol=1e-9,
)
return coupling_to_permutation(coupling)
SOLVERS = {
"grampa (current)": solve_grampa,
"umeyama": solve_umeyama,
"gromov-wasserstein": solve_gw,
"entropic GW annealed": solve_gw_annealed,
"BAPG GW": solve_bapg,
}
def main() -> None:
args = parse_args()
labels = args.labels or [p.split("/")[-1] for p in args.fields]
chosen = args.methods or list(SOLVERS)
device = torch.device(args.device)
rows = []
for path, label in zip(args.fields, labels):
state = torch.load(path, map_location="cpu", weights_only=False)
visual = standardise(state["visual_field"].double().numpy())
text = standardise(state["text_field"].double().numpy())
size = len(visual)
text_gpu = standardized(torch.from_numpy(text).to(device)).double()
visual_gpu = standardized(torch.from_numpy(visual).to(device)).double()
print(f"\n=== {label} (N={size}, chance={1/size:.4f})", flush=True)
for name in chosen:
solver = SOLVERS[name]
raw, refined, seconds, failures = [], [], [], []
for trial in range(args.trials):
generator = np.random.default_rng(trial)
hidden = generator.permutation(size)
shuffled = text[np.ix_(hidden, hidden)]
start = time.time()
try:
columns = solver(visual=visual, shuffled=shuffled)
except Exception as error:
failures.append(str(error)[:80])
continue
seconds.append(time.time() - start)
raw.append(float((hidden[columns] == np.arange(size)).mean()))
shuffled_gpu = standardized(
torch.from_numpy(shuffled).to(device)
).double()
final = fast_pair_descent(
shuffled_gpu, visual_gpu,
torch.from_numpy(np.ascontiguousarray(columns)).to(device),
args.iterations,
)
refined.append(
float((hidden[final.cpu().numpy()] == np.arange(size)).mean())
)
row = {
"field": label, "method": name,
"raw": float(np.mean(raw)) if raw else None,
"refined": float(np.mean(refined)) if refined else None,
"seconds": float(np.mean(seconds)) if seconds else None,
"failures": failures,
}
rows.append(row)
if raw:
print(" %-22s raw=%.4f refined=%.4f (%.1fs)"
% (name, row["raw"], row["refined"], row["seconds"]), flush=True)
else:
print(" %-22s FAILED: %s" % (name, failures[0] if failures else "?"),
flush=True)
write_json(args.output, {
"protocol": (
"Blind: the hidden permutation is generated per trial and read only "
"for scoring. Each solver is scored as returned and again after "
"exact steepest-descent refinement."
),
"rows": rows,
})
print(json.dumps({"done": True}))
if __name__ == "__main__":
main()
|
