summaryrefslogtreecommitdiff
path: root/worldalign/shared_dimension.py
blob: 2729a6b98b4da6b91e74b515653fb57cbb915cba (plain)
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
"""How many independent dimensions do the two modalities actually share?

The principal-angle count used earlier does not survive its own control. Adding
independent noise to both synthetic fields drives it from 21 to 15.7 while
recovery stays at 94.7%; suppressing one caption factor drives it to 15.0 and
recovery to 5.6%. Same correlation, same count, opposite outcome -- so the
count conflates two different things. Noise rotates eigenvectors and lowers the
measured overlap without narrowing the shared signal, which is still full-rank
underneath and still recoverable.

This measures the shared dimension in a way noise cannot fake, by asking how
many directions of agreement **generalise to scenes that took no part in
fitting them**. Scenes are split three ways. Anchors define a coordinate
system: a scene is described by its field row against the anchors, separately
in each modality. Canonical correlation is fitted on the second third and
evaluated on the last, so a direction counts only if the agreement it claims
holds up on scenes it never saw. Noise cannot survive that; genuine shared
structure can.

Nothing here uses the pairing beyond the alignment the two fields already
carry, and the count is calibrated against a scene-shuffled null.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from .common import write_json


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("--components", type=int, default=24)
    parser.add_argument("--ridge", type=float, default=1e-2)
    parser.add_argument("--threshold", type=float, default=0.3)
    parser.add_argument("--repeats", type=int, default=5)
    parser.add_argument("--output", default="artifacts/vg_5k/shared_dimension.json")
    return parser.parse_args()


def standardise(matrix: np.ndarray) -> np.ndarray:
    values = matrix[~np.eye(len(matrix), dtype=bool)]
    out = (matrix - values.mean()) / values.std()
    np.fill_diagonal(out, 0.0)
    return out


def reduce_columns(block: np.ndarray, components: int, basis=None):
    """PCA of anchor-similarity profiles, fitted once on the fit split."""
    if basis is None:
        centre = block.mean(0, keepdims=True)
        _, _, vectors = np.linalg.svd(block - centre, full_matrices=False)
        basis = (centre, vectors[:components].T)
    centre, directions = basis
    return (block - centre) @ directions, basis


def canonical(fit_a, fit_b, test_a, test_b, ridge: float) -> np.ndarray:
    """Held-out canonical correlations between two views of the same scenes."""
    fit_a = fit_a - fit_a.mean(0)
    fit_b = fit_b - fit_b.mean(0)
    n, dimension = fit_a.shape
    cov_aa = fit_a.T @ fit_a / n + ridge * np.eye(dimension)
    cov_bb = fit_b.T @ fit_b / n + ridge * np.eye(fit_b.shape[1])
    cov_ab = fit_a.T @ fit_b / n

    inv_a = np.linalg.inv(np.linalg.cholesky(cov_aa))
    inv_b = np.linalg.inv(np.linalg.cholesky(cov_bb))
    whitened = inv_a @ cov_ab @ inv_b.T
    left, _, right = np.linalg.svd(whitened, full_matrices=False)

    projection_a = (test_a - test_a.mean(0)) @ inv_a.T @ left
    projection_b = (test_b - test_b.mean(0)) @ inv_b.T @ right.T
    correlations = []
    for column in range(projection_a.shape[1]):
        first, second = projection_a[:, column], projection_b[:, column]
        denominator = first.std() * second.std()
        correlations.append(
            0.0 if denominator < 1e-12
            else float(np.mean((first - first.mean()) * (second - second.mean()))
                       / denominator)
        )
    return np.abs(np.array(correlations))


def analyse(path: str, label: str, args: argparse.Namespace) -> dict:
    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)
    mask = ~np.eye(size, dtype=bool)
    rho = float(np.corrcoef(visual[mask], text[mask])[0, 1])

    counts, spectra, nulls = [], [], []
    for repeat in range(args.repeats):
        generator = np.random.default_rng(repeat)
        order = generator.permutation(size)
        third = size // 3
        anchors, fit, test = order[:third], order[third : 2 * third], order[2 * third :]

        def views(rows, basis_v=None, basis_t=None):
            reduced_v, basis_v = reduce_columns(
                visual[np.ix_(rows, anchors)], args.components, basis_v
            )
            reduced_t, basis_t = reduce_columns(
                text[np.ix_(rows, anchors)], args.components, basis_t
            )
            return reduced_v, reduced_t, basis_v, basis_t

        fit_v, fit_t, basis_v, basis_t = views(fit)
        test_v, test_t, _, _ = views(test, basis_v, basis_t)
        correlations = canonical(fit_v, fit_t, test_v, test_t, args.ridge)
        counts.append(int((correlations > args.threshold).sum()))
        spectra.append(correlations)

        shuffled = text[np.ix_(generator.permutation(size), generator.permutation(size))]
        shuffled = standardise(shuffled)

        def shuffled_views(rows, basis=None):
            return reduce_columns(shuffled[np.ix_(rows, anchors)], args.components, basis)

        fit_s, basis_s = shuffled_views(fit)
        test_s, _ = shuffled_views(test, basis_s)
        nulls.append(
            int((canonical(fit_v, fit_s, test_v, test_s, args.ridge) > args.threshold).sum())
        )

    mean_spectrum = np.mean(np.stack(spectra), axis=0)
    return {
        "label": label,
        "correlation": rho,
        "shared_dimension": float(np.mean(counts)),
        "shuffled_null": float(np.mean(nulls)),
        "net_shared_dimension": float(np.mean(counts) - np.mean(nulls)),
        "held_out_canonical_top8": [round(float(v), 3) for v in mean_spectrum[:8]],
    }


def main() -> None:
    args = parse_args()
    labels = args.labels or [path.split("/")[-1] for path in args.fields]
    rows = []
    for path, label in zip(args.fields, labels):
        row = analyse(path, label, args)
        rows.append(row)
        print(
            f"{row['label']:<26} rho={row['correlation']:.3f}  "
            f"shared_dim={row['shared_dimension']:5.1f} "
            f"(null {row['shuffled_null']:4.1f}, net {row['net_shared_dimension']:5.1f})  "
            f"top={row['held_out_canonical_top8'][:5]}",
            flush=True,
        )
    write_json(args.output, {
        "protocol": (
            "Three-way scene split. Anchors define the coordinate system, "
            "canonical correlation is fitted on the second third and scored on "
            "the last, so a shared direction counts only if it generalises to "
            "scenes that took no part in fitting it. Calibrated against a "
            "scene-shuffled null."
        ),
        "components": args.components,
        "threshold": args.threshold,
        "rows": rows,
    })


if __name__ == "__main__":
    main()