diff options
Diffstat (limited to 'worldalign/natural_pipeline.py')
| -rw-r--r-- | worldalign/natural_pipeline.py | 78 |
1 files changed, 65 insertions, 13 deletions
diff --git a/worldalign/natural_pipeline.py b/worldalign/natural_pipeline.py index c30af4d..1a4f63e 100644 --- a/worldalign/natural_pipeline.py +++ b/worldalign/natural_pipeline.py @@ -41,9 +41,18 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--min-count", type=int, default=60) parser.add_argument("--max-vocabulary", type=int, default=600) parser.add_argument("--context-window", type=int, default=2) - parser.add_argument("--keep", type=int, default=8) + parser.add_argument("--keep", type=int, default=64) parser.add_argument("--shrinkage", type=float, default=0.05) + parser.add_argument( + "--weight-power", + type=float, + default=0.5, + help="Scale each kept direction by its discriminant eigenvalue to " + "this power. 0 reproduces the unweighted projection.", + ) parser.add_argument("--views", type=int, default=1) + parser.add_argument("--moment-degree", type=int, default=2, choices=[2, 3]) + parser.add_argument("--third-width", type=int, default=12) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--output", default="artifacts/vg_5k/natural_pipeline.pt") return parser.parse_args() @@ -77,10 +86,43 @@ def word_vectors( return {word: embedding[row] for row, word in enumerate(vocabulary)} +def moment_field_degree( + sets: list[torch.Tensor], degree: int, third_width: int +) -> torch.Tensor: + """Set kernel carrying moments up to `degree`. + + A kernel carrying third moments spans more directions than one carrying + two, so raising the degree raises the field's rank by construction -- + which is the quantity the rank experiment showed to govern recovery. The + third moment is taken over a narrower basis because its size is cubic; + the first two are kept at full width. + """ + if degree < 3: + return moment_field(sets) + phis = [] + for members in sets: + first = members.mean(0) + second = (members[:, :, None] * members[:, None, :]).mean(0).flatten() + narrow = members[:, :third_width] + outer = narrow[:, :, None] * narrow[:, None, :] + third = (outer[:, :, :, None] * narrow[:, None, None, :]).mean(0).flatten() + phi = torch.cat([first, second, third]) + phis.append(phi / phi.norm().clamp_min(1e-9)) + stacked = torch.stack(phis) + return stacked @ stacked.T + + def content_directions( sets: list[np.ndarray], shrinkage: float -) -> tuple[np.ndarray, np.ndarray]: - """Directions separating scenes more than they separate parts.""" +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Directions separating scenes more than they separate parts. + + Returns the discriminant eigenvalues alongside the basis. Scaling each + direction by its own eigenvalue lets a wide basis be kept without the + weakly separating directions drowning the strong ones, which is what + makes a wide keep beat a narrow one -- and a wide keep is what raises + the rank of the resulting field. + """ means = np.stack([item.mean(0) for item in sets]) centre = means.mean(0) within = np.concatenate([item - item.mean(0, keepdims=True) for item in sets]) @@ -91,7 +133,8 @@ def content_directions( values, vectors = eigh( scatter_between, scatter_within + shrinkage * trace * np.eye(len(scatter_within)) ) - return centre, vectors[:, np.argsort(values)[::-1]] + order = np.argsort(values)[::-1] + return centre, vectors[:, order], np.maximum(values[order], 0.0) def main() -> None: @@ -144,8 +187,8 @@ def main() -> None: language_fit = [item for item in language_fit if item is not None and len(item) > 1] vision_fit = [vision_state(p["vision_node_id"], 0) for p in fit] vision_fit = [item for item in vision_fit if item is not None and len(item) > 1] - language_centre, language_basis = content_directions(language_fit, args.shrinkage) - vision_centre, vision_basis = content_directions(vision_fit, args.shrinkage) + language_fitted = content_directions(language_fit, args.shrinkage) + vision_fitted = content_directions(vision_fit, args.shrinkage) evaluated = [ pair for pair in pairs[: args.samples] @@ -153,25 +196,32 @@ def main() -> None: ] keep = args.keep - def project(raw: np.ndarray, centre: np.ndarray, basis: np.ndarray) -> torch.Tensor: + def project(raw: np.ndarray, fitted: tuple) -> torch.Tensor: + centre, basis, values = fitted width = min(keep, basis.shape[1]) + scale = values[:width] ** args.weight_power if args.weight_power else 1.0 return F.normalize( - torch.tensor((raw - centre) @ basis[:, :width], dtype=torch.float32), dim=-1 + torch.tensor( + (raw - centre) @ basis[:, :width] * scale, dtype=torch.float32 + ), dim=-1 ) text_sets = [ - project(language_state(p["text_node_id"]), language_centre, language_basis) - for p in evaluated + project(language_state(p["text_node_id"]), language_fitted) for p in evaluated ] view_fields = [] for view in range(views): sets = [ - project(vision_state(p["vision_node_id"], view), vision_centre, vision_basis) + project(vision_state(p["vision_node_id"], view), vision_fitted) for p in evaluated ] - view_fields.append(moment_field(sets)) + view_fields.append( + moment_field_degree(sets, args.moment_degree, args.third_width) + ) visual_field = torch.stack(view_fields).mean(0) - text_field = moment_field(text_sets) + text_field = moment_field_degree( + text_sets, args.moment_degree, args.third_width + ) mask = ~np.eye(len(evaluated), dtype=bool) correlation = float( @@ -193,6 +243,8 @@ def main() -> None: "samples": len(evaluated), "views": views, "kept_directions": keep, + "weight_power": args.weight_power, + "moment_degree": args.moment_degree, "field_correlation_at_truth": correlation, "recovery_threshold": 0.9, } |