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
|
"""Close the loop: recovered pairs to a map to generated descriptions.
Matching is transductive -- it aligns one fixed population. A usable
system needs a map, so the recovered assignment is treated as pseudo-pair
supervision for a small vision-to-text-state regressor, and the map is
then applied to scenes that took no part in the matching. Descriptions
are read out through the frozen text tower.
Two controls decide whether the output is image-conditioned: the same
pipeline with the recovered assignment replaced by a random one, and the
same map applied to a shuffled image. Hidden pairs score; nothing here
trains on them.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
from .common import read_json, seed_everything, write_json
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--data-dir", default="artifacts/synth_v1")
parser.add_argument("--states", required=True,
help=".pt with vision_states, text_states, rows for "
"the matched population and the held-out split.")
parser.add_argument("--assignment", required=True,
help=".pt with the recovered permutation and truth.")
parser.add_argument("--ridge", type=float, default=1.0)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--output", default="artifacts/synth_v1/decode.json")
return parser.parse_args()
def fit_map(source: np.ndarray, target: np.ndarray, ridge: float) -> np.ndarray:
source = np.concatenate([source, np.ones((len(source), 1))], axis=1)
gram = source.T @ source + ridge * np.eye(source.shape[1])
return np.linalg.solve(gram, source.T @ target)
def apply_map(weights: np.ndarray, source: np.ndarray) -> np.ndarray:
source = np.concatenate([source, np.ones((len(source), 1))], axis=1)
return source @ weights
def nearest_caption(
predicted: np.ndarray, bank: np.ndarray
) -> np.ndarray:
predicted = predicted / np.linalg.norm(predicted, axis=1, keepdims=True).clip(1e-9)
bank = bank / np.linalg.norm(bank, axis=1, keepdims=True).clip(1e-9)
return (predicted @ bank.T).argmax(1)
def factor_scores(
predicted_rows: list[int], truth_rows: list[int], scenes: list[dict]
) -> dict:
"""Do the retrieved descriptions state the right world facts?"""
colour_f1, count_ok, group_ok = [], [], []
for predicted, truth in zip(predicted_rows, truth_rows):
p, t = scenes[predicted], scenes[truth]
pc = {g["color"] for g in p["groups"]}
tc = {g["color"] for g in t["groups"]}
overlap = len(pc & tc)
precision = overlap / max(len(pc), 1)
recall = overlap / max(len(tc), 1)
colour_f1.append(
0.0 if precision + recall == 0 else 2 * precision * recall / (precision + recall)
)
count_ok.append(
sorted(g["count"] for g in p["groups"])
== sorted(g["count"] for g in t["groups"])
)
group_ok.append(len(p["groups"]) == len(t["groups"]))
return {
"colour_set_f1": float(np.mean(colour_f1)),
"count_multiset_exact": float(np.mean(count_ok)),
"group_count_exact": float(np.mean(group_ok)),
}
def main() -> None:
args = parse_args()
seed_everything(args.seed)
scenes = read_json(Path(args.data_dir, "scenes.private.json"))["scenes"]
states = torch.load(args.states, map_location="cpu", weights_only=False)
recovered = torch.load(args.assignment, map_location="cpu", weights_only=False)
match_vision = states["match_vision"].double().numpy()
match_text = states["match_text"].double().numpy()
held_vision = states["held_vision"].double().numpy()
held_text = states["held_text"].double().numpy()
held_rows = states["held_rows"]
assignment = np.asarray(recovered["assignment"])
generator = np.random.default_rng(args.seed)
random_assignment = generator.permutation(len(assignment))
report = {
"protocol": (
"The recovered assignment supplies pseudo-pairs for a ridge "
"map from vision states to text states; the map is applied to "
"held-out scenes that took no part in matching. Random "
"assignment and shuffled-image controls bound the claim."
),
"matched_population": len(assignment),
"held_out": len(held_rows),
"recovery_accuracy": float(
(assignment == np.arange(len(assignment))).mean()
),
"conditions": {},
}
for label, permutation in (
("recovered_pairs", assignment),
("random_pairs", random_assignment),
):
weights = fit_map(match_vision, match_text[permutation], args.ridge)
predicted = apply_map(weights, held_vision)
retrieved = nearest_caption(predicted, held_text)
exact = float((retrieved == np.arange(len(held_rows))).mean())
entry = {
"held_out_retrieval_exact": exact,
"chance": 1.0 / len(held_rows),
**factor_scores(
[held_rows[i] for i in retrieved], held_rows, scenes
),
}
if label == "recovered_pairs":
shuffled = generator.permutation(len(held_vision))
predicted_shuffled = apply_map(weights, held_vision[shuffled])
retrieved_shuffled = nearest_caption(predicted_shuffled, held_text)
entry["shuffled_image_control"] = factor_scores(
[held_rows[i] for i in retrieved_shuffled], held_rows, scenes
)
report["conditions"][label] = entry
print(json.dumps({label: entry}))
write_json(args.output, report)
print(f"Wrote {args.output}")
if __name__ == "__main__":
main()
|
