""" Equalized Signed-Hash Subspace Experiment Runner for Heavy Task Autoresearch. Protocol Version: HEAVY-PSO-AUTORESEARCH 1.0.0 Runs matched equalized signed-hash subspace PSO experiments across four candidate latent ratios (0.5, 0.25, 0.125, 0.03125) and four heavy workloads: - mnist_compact (Base geometry: G6) - mnist_wide (Base geometry: G5) - fashion_compact (Base geometry: G6) - fashion_wide (Base geometry: G5) Configurations are matched to confirmation runs (12 particles, 80 epochs, fixed 10k, seeds 101-103). Official test splits are NEVER loaded or evaluated (official_test_evaluations = 0). """ from __future__ import annotations import argparse import json import hashlib import math import sys import time from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Union import numpy as np import torch import torch.nn as nn # Ensure test directory and repo root are in Python path REPO_ROOT = Path(__file__).resolve().parent.parent if Path(__file__).resolve().parent.name != "PSO" else Path(__file__).resolve().parent TEST_DIR = REPO_ROOT / "test" if str(TEST_DIR) not in sys.path: sys.path.insert(0, str(TEST_DIR)) if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from benchmark_suite import ( calc_stats, compute_model_fingerprint, get_hardware_provenance, resolve_execution_device, save_json_atomic, ) from deep_pso_v6 import ( V6GeometryConfig, V6LatentTransform, compute_equalized_subspace_radius, get_v6_geometry_table, run_v6_pso, artifact_safe_run, ) from heavy_task_feasibility import ( WORKLOADS, create_model, prepare_heavy_task_data, ) # Protocol Constant AUTORESEARCH_PROTOCOL_VERSION = "HEAVY-PSO-AUTORESEARCH 1.0.0" # Defaults DEFAULT_RATIOS = (1.0, 0.5, 0.25, 0.125, 0.03125) DEFAULT_PARTICLES = 12 DEFAULT_EPOCHS = 80 DEFAULT_SUBSET_SIZE = 10000 DEFAULT_SEEDS = (101, 102, 103) DEFAULT_SPLIT_SEED = 20260902 DEFAULT_GEOMETRY_POLICY = "recovered" DEFAULT_PROJECTION_SCOPE = "global" PROJECTION_SCOPES = ("global", "tensor_local", "balanced_global", "two_hash_global", "largest_tensor_hash", "largest_tensor_row_hash", "adjacent_pair", "adjacent_difference") DEFAULT_PROJECTION_SEED_MODE = "coupled" PROJECTION_SEED_MODES = ("coupled", "fixed", "explicit") DEFAULT_GEOMETRY_MULTIPLIER = 1.0 def validate_geometry_multiplier(geometry_multiplier: Any) -> None: """ Validates that geometry_multiplier is a finite positive float (> 0). Raises ValueError on non-numeric, non-finite, or non-positive values. """ if ( not isinstance(geometry_multiplier, (int, float)) or isinstance(geometry_multiplier, bool) or not math.isfinite(geometry_multiplier) or geometry_multiplier <= 0 ): raise ValueError( f"geometry_multiplier must be a finite positive float (> 0), got {geometry_multiplier!r}" ) def parse_projection_seed_arg(val: Any) -> Optional[Union[int, Dict[str, int]]]: """ Parses projection_seed parameter or CLI arg into None, int, or Dict[str, int]. Accepts integer values, integer strings, JSON dict strings, or comma/colon key-value strings. """ if val is None or val == "" or val == "None": return None if isinstance(val, (int, dict)): return val if isinstance(val, str): val_str = val.strip() if not val_str: return None try: return int(val_str) except ValueError: pass if val_str.startswith("{") and val_str.endswith("}"): try: parsed = json.loads(val_str) if isinstance(parsed, dict): return {str(k): int(v) for k, v in parsed.items()} except (json.JSONDecodeError, ValueError, TypeError) as exc: raise ValueError(f"Failed to parse projection_seed JSON dict string '{val}': {exc}") if ":" in val_str or "=" in val_str: res = {} for item in val_str.replace(";", ",").split(","): item = item.strip() if not item: continue if ":" in item: k, v = item.split(":", 1) elif "=" in item: k, v = item.split("=", 1) else: raise ValueError(f"Invalid key-value projection_seed string item '{item}'") res[k.strip()] = int(v.strip()) if res: return res raise ValueError(f"Cannot parse projection_seed value: {val!r}") def validate_projection_seed_config( projection_seed_mode: str, projection_seed: Optional[Union[int, Dict[str, int]]], ) -> None: """ Validates projection_seed_mode and projection_seed invariants. Raises ValueError on invalid modes, missing explicit seeds, out-of-range explicit seeds, or seeds supplied to non-explicit modes. """ if projection_seed_mode not in PROJECTION_SEED_MODES: raise ValueError( f"Invalid projection_seed_mode '{projection_seed_mode}'. Must be one of {list(PROJECTION_SEED_MODES)}" ) if projection_seed_mode == "explicit": if projection_seed is None: raise ValueError( "projection_seed must be provided when projection_seed_mode is 'explicit'" ) if isinstance(projection_seed, int) and not isinstance(projection_seed, bool): if not (0 <= projection_seed < (2**31 - 1)): raise ValueError( f"projection_seed must be a non-negative integer < 2**31-1, got {projection_seed!r}" ) elif isinstance(projection_seed, dict): if not projection_seed: raise ValueError("projection_seed dictionary cannot be empty") expected_keys = set(WORKLOADS.keys()) provided_keys = set(projection_seed.keys()) missing_keys = expected_keys - provided_keys unknown_keys = provided_keys - expected_keys if missing_keys or unknown_keys: details = [] if missing_keys: details.append(f"missing required workload key(s) {sorted(missing_keys)}") if unknown_keys: details.append(f"Unknown workload_id key(s) {sorted(unknown_keys)}") raise ValueError( f"projection_seed dictionary must contain exactly WORKLOADS keys ({sorted(expected_keys)}); " + ", ".join(details) ) for k, v in projection_seed.items(): if ( not isinstance(v, int) or isinstance(v, bool) or not (0 <= v < (2**31 - 1)) ): raise ValueError( f"projection_seed for workload '{k}' must be a non-negative integer < 2**31-1, got {v!r}" ) else: raise ValueError( f"projection_seed must be a non-negative integer < 2**31-1 or a " f"Dict[str, int], got {projection_seed!r}" ) else: if projection_seed is not None: raise ValueError( f"projection_seed can only be provided when projection_seed_mode is 'explicit', got mode='{projection_seed_mode}' and projection_seed={projection_seed!r}" ) def parse_projection_scope_arg(val: Any) -> Union[str, Dict[str, str]]: """ Parses projection_scope parameter or CLI arg into a string or Dict[str, str]. Accepts valid scope strings, JSON dict strings, or comma/colon key-value strings. """ if val is None or val == "": return DEFAULT_PROJECTION_SCOPE if isinstance(val, dict): return {str(k): str(v) for k, v in val.items()} if isinstance(val, str): val_str = val.strip() if not val_str: return DEFAULT_PROJECTION_SCOPE if val_str in PROJECTION_SCOPES: return val_str if val_str.startswith("{") and val_str.endswith("}"): try: parsed = json.loads(val_str) if isinstance(parsed, dict): return {str(k): str(v) for k, v in parsed.items()} except Exception as exc: raise ValueError(f"Failed to parse projection_scope JSON dict string '{val}': {exc}") if ":" in val_str or "=" in val_str: res = {} for item in val_str.replace(";", ",").split(","): item = item.strip() if not item: continue if ":" in item: k, v = item.split(":", 1) elif "=" in item: k, v = item.split("=", 1) else: raise ValueError(f"Invalid key-value projection_scope string item '{item}'") res[k.strip()] = v.strip() if res: return res return val_str raise ValueError(f"Cannot parse projection_scope value: {val!r}") def validate_projection_scope_config(projection_scope: Any) -> None: """ Validates projection_scope string or dictionary config. Raises ValueError if string is not in PROJECTION_SCOPES, or if dict keys do not match WORKLOADS exactly or values are not in PROJECTION_SCOPES. """ if isinstance(projection_scope, str): if projection_scope not in PROJECTION_SCOPES: raise ValueError( f"Invalid projection_scope '{projection_scope}'. Must be one of {list(PROJECTION_SCOPES)}" ) elif isinstance(projection_scope, dict): if not projection_scope: raise ValueError("projection_scope dictionary cannot be empty") if set(projection_scope.keys()) != set(WORKLOADS.keys()): missing = sorted(list(set(WORKLOADS.keys()) - set(projection_scope.keys()))) extra = sorted(list(set(projection_scope.keys()) - set(WORKLOADS.keys()))) details = [] if missing: details.append(f"missing keys {missing}") if extra: details.append(f"unknown keys {extra}") raise ValueError( f"projection_scope dictionary must contain exact workload keys {sorted(list(WORKLOADS.keys()))}, got {', '.join(details)}" ) for k, v in projection_scope.items(): if v not in PROJECTION_SCOPES: raise ValueError( f"Invalid projection_scope '{v}' for workload '{k}'. Must be one of {list(PROJECTION_SCOPES)}" ) else: raise ValueError( f"projection_scope must be a string or Dict[str, str], got {projection_scope!r}" ) def get_effective_projection_scope( projection_scope: Union[str, Dict[str, str]], workload_id: str, ) -> str: """ Returns the effective projection scope string for a given workload. """ if isinstance(projection_scope, dict): if workload_id not in projection_scope: raise ValueError(f"Missing workload_id '{workload_id}' in projection_scope dict") return str(projection_scope[workload_id]) return str(projection_scope) def allocate_tensor_latent_dims( param_numels: Sequence[int], aggregate_latent_dim: int, ) -> List[int]: """ Allocates aggregate_latent_dim across parameter tensors deterministically, proportionally to tensor numel, with at least one coordinate per tensor, no tensor exceeding numel, and exact sum aggregate_latent_dim. """ total_dim = sum(param_numels) if not (0 < aggregate_latent_dim <= total_dim): raise ValueError( f"aggregate_latent_dim must be in (0, {total_dim}], got {aggregate_latent_dim}" ) num_tensors = len(param_numels) if aggregate_latent_dim < num_tensors: raise ValueError( f"aggregate_latent_dim ({aggregate_latent_dim}) must be at least number of parameter tensors ({num_tensors})" ) if aggregate_latent_dim == total_dim: return list(param_numels) quotas = [aggregate_latent_dim * n / total_dim for n in param_numels] allocs = [max(1, min(n, int(math.floor(q)))) for n, q in zip(param_numels, quotas)] current_sum = sum(allocs) if current_sum < aggregate_latent_dim: deficit = aggregate_latent_dim - current_sum candidates = [i for i in range(num_tensors) if allocs[i] < param_numels[i]] candidates.sort( key=lambda i: (quotas[i] - math.floor(quotas[i]), param_numels[i], -i), reverse=True, ) for i in candidates[:deficit]: allocs[i] += 1 elif current_sum > aggregate_latent_dim: surplus = current_sum - aggregate_latent_dim candidates = [i for i in range(num_tensors) if allocs[i] > 1] candidates.sort( key=lambda i: (quotas[i] - math.floor(quotas[i]), param_numels[i], -i), reverse=False, ) for i in candidates[:surplus]: allocs[i] -= 1 assert sum(allocs) == aggregate_latent_dim, ( f"Allocation sum {sum(allocs)} does not match aggregate_latent_dim {aggregate_latent_dim}" ) assert all(1 <= a <= n for a, n in zip(allocs, param_numels)), ( f"Allocation bounds violated: {allocs} vs numels {param_numels}" ) return allocs class TensorLocalLatentTransform(V6LatentTransform): """ Subclass of V6LatentTransform that isolates signed-hash coordinates within each model parameter tensor, mapping each tensor's parameters only to its assigned contiguous latent slice. """ def __init__( self, base_model: nn.Module, geom_config: V6GeometryConfig, device: torch.device, ): super().__init__(base_model, geom_config, device) if not self.is_full: self.tensor_latent_dims = allocate_tensor_latent_dims( self.param_numels, self.latent_dim ) k_indices = np.zeros(self.total_dim, dtype=np.int64) h2_signs = np.zeros(self.total_dim, dtype=np.float32) seed_offset = ( geom_config.projection_seed if geom_config.projection_seed is not None else 0 ) j_offset = 0 l_offset = 0 for numel, d_m in zip(self.param_numels, self.tensor_latent_dims): j_local = np.arange(numel, dtype=np.int64) j_global = j_offset + j_local h1 = ((j_global + 1 + seed_offset) * 2654435761) % (2**32) k_local = h1 % d_m k_indices[j_global] = l_offset + k_local h2 = ((j_global + 1 + seed_offset) * 1597334677) % (2**32) h2_signs[j_global] = np.where((h2 % 2) == 0, 1.0, -1.0) j_offset += numel l_offset += d_m bin_counts = np.bincount(k_indices, minlength=self.latent_dim) count_per_j = bin_counts[k_indices] scale_per_j = 1.0 / np.sqrt(np.maximum(count_per_j, 1)) combined_weights = h2_signs * scale_per_j self.k_indices = torch.tensor(k_indices, dtype=torch.long, device=device) self.weights = torch.tensor(combined_weights, dtype=torch.float32, device=device) else: self.tensor_latent_dims = list(self.param_numels) class BalancedGlobalLatentTransform(V6LatentTransform): """ Subclass of V6LatentTransform that enforces balanced parameter occupancy across global latent coordinate buckets using a deterministic permutation and sign assignment. """ def __init__( self, base_model: nn.Module, geom_config: V6GeometryConfig, device: torch.device, ): super().__init__(base_model, geom_config, device) if not self.is_full: seed_offset = ( geom_config.projection_seed if geom_config.projection_seed is not None else 0 ) rng = np.random.RandomState(seed_offset) perm = rng.permutation(self.total_dim) k_indices = np.zeros(self.total_dim, dtype=np.int64) k_indices[perm] = np.arange(self.total_dim, dtype=np.int64) % self.latent_dim h2_signs = rng.choice(np.array([1.0, -1.0], dtype=np.float32), size=self.total_dim) bin_counts = np.bincount(k_indices, minlength=self.latent_dim) count_per_j = bin_counts[k_indices] scale_per_j = 1.0 / np.sqrt(np.maximum(count_per_j, 1)) combined_weights = h2_signs * scale_per_j self.k_indices = torch.tensor(k_indices, dtype=torch.long, device=device) self.weights = torch.tensor(combined_weights, dtype=torch.float32, device=device) class TwoHashGlobalLatentTransform(V6LatentTransform): """ Subclass of V6LatentTransform that maps each parameter to two distinct global latent coordinate buckets using independent deterministic signed-hash mappings, normalized by per-bucket 1/sqrt(count) weights and decoded as (term1 + term2)/sqrt(2). """ def __init__( self, base_model: nn.Module, geom_config: V6GeometryConfig, device: torch.device, ): super().__init__(base_model, geom_config, device) if not self.is_full: j_indices = np.arange(self.total_dim, dtype=np.int64) seed_offset = ( geom_config.projection_seed if geom_config.projection_seed is not None else 0 ) # Hash Map 1 (Primary signed hash projection) h1_1 = ((j_indices + 1 + seed_offset) * 2654435761) % (2**32) k1_indices = h1_1 % self.latent_dim h1_2 = ((j_indices + 1 + seed_offset) * 1597334677) % (2**32) signs1 = np.where((h1_2 % 2) == 0, 1.0, -1.0) bin_counts1 = np.bincount(k1_indices, minlength=self.latent_dim) count_per_j1 = bin_counts1[k1_indices] scale_per_j1 = 1.0 / np.sqrt(np.maximum(count_per_j1, 1)) combined_weights1 = signs1 * scale_per_j1 # Hash Map 2 (Secondary independent signed hash projection) h2_1 = ((j_indices + 1 + seed_offset) * 2246822519) % (2**32) if self.latent_dim > 1: offset = 1 + (h2_1 % (self.latent_dim - 1)) k2_indices = (k1_indices + offset) % self.latent_dim else: k2_indices = np.zeros(self.total_dim, dtype=np.int64) h2_2 = ((j_indices + 1 + seed_offset) * 3266489917) % (2**32) signs2 = np.where(((h2_2 >> 16) % 2) == 0, 1.0, -1.0) bin_counts2 = np.bincount(k2_indices, minlength=self.latent_dim) count_per_j2 = bin_counts2[k2_indices] scale_per_j2 = 1.0 / np.sqrt(np.maximum(count_per_j2, 1)) combined_weights2 = signs2 * scale_per_j2 self.k1_indices = torch.tensor(k1_indices, dtype=torch.long, device=device) self.weights1 = torch.tensor(combined_weights1, dtype=torch.float32, device=device) self.k2_indices = torch.tensor(k2_indices, dtype=torch.long, device=device) self.weights2 = torch.tensor(combined_weights2, dtype=torch.float32, device=device) # Backward compatibility aliases self.k_indices = self.k1_indices self.weights = self.weights1 def decode(self, Z: torch.Tensor) -> torch.Tensor: """ Transforms latent batch Z (N, d) into full parameter batch (N, D). When not full, term1 = Z[:, k1] * weights1, term2 = Z[:, k2] * weights2, delta = (term1 + term2) / sqrt(2). theta = base_vec + scale_vec * delta """ if self.is_full: delta = Z else: term1 = Z[:, self.k1_indices] * self.weights1 term2 = Z[:, self.k2_indices] * self.weights2 delta = (term1 + term2) / math.sqrt(2.0) return self.base_vec + self.scale_vec * delta class LargestTensorHashLatentTransform(V6LatentTransform): """ Subclass of V6LatentTransform that isolates the single largest parameter tensor by numel (with stable first-index tie break) for signed-hash projection, while assigning every parameter in all other tensors a unique direct latent coordinate with weight 1.0. """ def __init__( self, base_model: nn.Module, geom_config: V6GeometryConfig, device: torch.device, ): super().__init__(base_model, geom_config, device) if not self.is_full: largest_idx = int(np.argmax(self.param_numels)) protected_dim = sum( numel for i, numel in enumerate(self.param_numels) if i != largest_idx ) residual_dim = self.latent_dim - protected_dim if residual_dim < 1: raise ValueError( f"latent_dim ({self.latent_dim}) must be greater than protected non-largest parameter dimension ({protected_dim}) " f"to provide at least 1 residual latent coordinate for the largest tensor (index {largest_idx}, numel {self.param_numels[largest_idx]})" ) k_indices = np.zeros(self.total_dim, dtype=np.int64) weights = np.zeros(self.total_dim, dtype=np.float32) seed_offset = ( geom_config.projection_seed if geom_config.projection_seed is not None else 0 ) j_offset = 0 direct_coord = 0 for i, numel in enumerate(self.param_numels): j_indices_tensor = np.arange(j_offset, j_offset + numel, dtype=np.int64) if i != largest_idx: k_indices[j_indices_tensor] = direct_coord + np.arange(numel, dtype=np.int64) weights[j_indices_tensor] = 1.0 direct_coord += numel else: h1 = ((j_indices_tensor + 1 + seed_offset) * 2654435761) % (2**32) k_rel = h1 % residual_dim k_indices[j_indices_tensor] = protected_dim + k_rel h2 = ((j_indices_tensor + 1 + seed_offset) * 1597334677) % (2**32) h2_signs = np.where((h2 % 2) == 0, 1.0, -1.0) weights[j_indices_tensor] = h2_signs j_offset += numel bin_counts = np.bincount(k_indices, minlength=self.latent_dim) count_per_j = bin_counts[k_indices] scale_per_j = 1.0 / np.sqrt(np.maximum(count_per_j, 1)) combined_weights = weights * scale_per_j self.k_indices = torch.tensor(k_indices, dtype=torch.long, device=device) self.weights = torch.tensor(combined_weights, dtype=torch.float32, device=device) class LargestTensorRowHashLatentTransform(V6LatentTransform): """ Subclass of V6LatentTransform that isolates the single largest parameter tensor by numel (with stable first-index tie break) for row-partitioned signed-hash projection across its first dimension (rows), while assigning every parameter in all other tensors a unique direct latent coordinate with weight 1.0. """ def __init__( self, base_model: nn.Module, geom_config: V6GeometryConfig, device: torch.device, ): super().__init__(base_model, geom_config, device) largest_idx = int(np.argmax(self.param_numels)) shape = self.param_shapes[largest_idx] num_rows = shape[0] if len(shape) >= 2 else 1 largest_numel = self.param_numels[largest_idx] if len(shape) >= 2: elements_per_row = largest_numel // num_rows row_numels = [elements_per_row] * num_rows else: row_numels = [largest_numel] if not self.is_full: protected_dim = sum( numel for i, numel in enumerate(self.param_numels) if i != largest_idx ) residual_dim = self.latent_dim - protected_dim if residual_dim < num_rows: raise ValueError( f"latent_dim ({self.latent_dim}) must be at least protected dimension ({protected_dim}) " f"+ number of rows ({num_rows}) for largest tensor row hashing, got residual_dim {residual_dim}" ) row_latent_dims = allocate_tensor_latent_dims(row_numels, residual_dim) k_indices = np.zeros(self.total_dim, dtype=np.int64) weights = np.zeros(self.total_dim, dtype=np.float32) seed_offset = ( geom_config.projection_seed if geom_config.projection_seed is not None else 0 ) j_offset = 0 direct_coord = 0 for i, numel in enumerate(self.param_numels): if i != largest_idx: j_indices_tensor = np.arange(j_offset, j_offset + numel, dtype=np.int64) k_indices[j_indices_tensor] = direct_coord + np.arange(numel, dtype=np.int64) weights[j_indices_tensor] = 1.0 direct_coord += numel j_offset += numel else: row_slice_start = protected_dim for r, (r_numel, r_dim) in enumerate(zip(row_numels, row_latent_dims)): j_indices_row = np.arange(j_offset, j_offset + r_numel, dtype=np.int64) h1 = ((j_indices_row + 1 + seed_offset) * 2654435761) % (2**32) k_rel = h1 % r_dim k_indices[j_indices_row] = row_slice_start + k_rel h2 = ((j_indices_row + 1 + seed_offset) * 1597334677) % (2**32) h2_signs = np.where((h2 % 2) == 0, 1.0, -1.0) weights[j_indices_row] = h2_signs j_offset += r_numel row_slice_start += r_dim bin_counts = np.bincount(k_indices, minlength=self.latent_dim) count_per_j = bin_counts[k_indices] scale_per_j = 1.0 / np.sqrt(np.maximum(count_per_j, 1)) combined_weights = weights * scale_per_j self.k_indices = torch.tensor(k_indices, dtype=torch.long, device=device) self.weights = torch.tensor(combined_weights, dtype=torch.float32, device=device) self.row_latent_dims = row_latent_dims else: self.row_latent_dims = row_numels class AdjacentPairLatentTransform(V6LatentTransform): """ Subclass of V6LatentTransform that maps each parameter tensor independently by assigning consecutive pairs of parameters to one unique latent coordinate with weights 1/sqrt(2), and a final unpaired parameter (if any) to weight 1.0, concatenating tensor coordinate ranges without sharing. """ def __init__( self, base_model: nn.Module, geom_config: V6GeometryConfig, device: torch.device, ): super().__init__(base_model, geom_config, device) if not self.is_full: required_dim = sum(math.ceil(n / 2) for n in self.param_numels) if self.latent_dim != required_dim: raise ValueError( f"AdjacentPairLatentTransform requires latent_dim == sum(ceil(numel_i/2)) = {required_dim}, got {self.latent_dim}" ) k_indices = np.zeros(self.total_dim, dtype=np.int64) weights = np.zeros(self.total_dim, dtype=np.float32) inv_sqrt2 = 1.0 / math.sqrt(2.0) j_offset = 0 l_offset = 0 self.tensor_latent_dims = [] for numel in self.param_numels: d_m = math.ceil(numel / 2) self.tensor_latent_dims.append(d_m) for p in range(numel): j = j_offset + p k_local = p // 2 k_indices[j] = l_offset + k_local if p % 2 == 0 and p == numel - 1: weights[j] = 1.0 else: weights[j] = inv_sqrt2 j_offset += numel l_offset += d_m self.k_indices = torch.tensor(k_indices, dtype=torch.long, device=device) self.weights = torch.tensor(weights, dtype=torch.float32, device=device) else: self.tensor_latent_dims = list(self.param_numels) class AdjacentDifferenceLatentTransform(AdjacentPairLatentTransform): """ Subclass of AdjacentPairLatentTransform that maps each parameter tensor independently by assigning consecutive pairs of parameters to one unique latent coordinate with opposite weights (+1/sqrt(2), -1/sqrt(2)), and a final unpaired parameter (if any) to weight 1.0, concatenating tensor coordinate ranges without sharing. """ def __init__( self, base_model: nn.Module, geom_config: V6GeometryConfig, device: torch.device, ): try: super().__init__(base_model, geom_config, device) except ValueError as e: raise ValueError( str(e).replace("AdjacentPairLatentTransform", "AdjacentDifferenceLatentTransform") ) from None if not self.is_full: weights = self.weights.clone() j_offset = 0 for numel in self.param_numels: for p in range(1, numel, 2): weights[j_offset + p] = -weights[j_offset + p] j_offset += numel self.weights = weights def format_ratio_id( ratio: float, geometry_policy: str = DEFAULT_GEOMETRY_POLICY, projection_scope: Union[str, Dict[str, str]] = DEFAULT_PROJECTION_SCOPE, projection_seed_mode: str = DEFAULT_PROJECTION_SEED_MODE, projection_seed: Optional[Union[int, Dict[str, int]]] = None, geometry_multiplier: float = DEFAULT_GEOMETRY_MULTIPLIER, ) -> str: validate_geometry_multiplier(geometry_multiplier) validate_projection_seed_config(projection_seed_mode, projection_seed) validate_projection_scope_config(projection_scope) r_str = f"{ratio:g}" base_id = f"aligned_r{r_str}" if geometry_policy == "baseline_aligned" else f"r{r_str}" if isinstance(projection_scope, dict): unique_scopes = set(projection_scope.values()) if len(unique_scopes) == 1: eff_scope = next(iter(unique_scopes)) if eff_scope == "tensor_local": base_id = f"local_{base_id}" elif eff_scope == "balanced_global": base_id = f"balanced_{base_id}" elif eff_scope == "two_hash_global": base_id = f"two_hash_{base_id}" elif eff_scope == "largest_tensor_hash": base_id = f"largest_tensor_hash_{base_id}" elif eff_scope == "largest_tensor_row_hash": base_id = f"largest_tensor_row_hash_{base_id}" elif eff_scope == "adjacent_pair": base_id = f"adjacent_pair_{base_id}" elif eff_scope == "adjacent_difference": base_id = f"adjacent_difference_{base_id}" else: base_id = f"mixed_{base_id}" elif projection_scope == "tensor_local": base_id = f"local_{base_id}" elif projection_scope == "balanced_global": base_id = f"balanced_{base_id}" elif projection_scope == "two_hash_global": base_id = f"two_hash_{base_id}" elif projection_scope == "largest_tensor_hash": base_id = f"largest_tensor_hash_{base_id}" elif projection_scope == "largest_tensor_row_hash": base_id = f"largest_tensor_row_hash_{base_id}" elif projection_scope == "adjacent_pair": base_id = f"adjacent_pair_{base_id}" elif projection_scope == "adjacent_difference": base_id = f"adjacent_difference_{base_id}" if projection_seed_mode == "fixed": base_id = f"fixed_{base_id}" elif projection_seed_mode == "explicit": if isinstance(projection_seed, dict): base_id = f"pexplicit_{base_id}" else: base_id = f"p{projection_seed}_{base_id}" if float(geometry_multiplier) != 1.0: g_str = f"{geometry_multiplier:g}" base_id = f"g{g_str}_{base_id}" return base_id # Geometry Policies Assignment GEOMETRY_POLICIES: Dict[str, Dict[str, str]] = { "recovered": { "mnist_compact": "G6", "mnist_wide": "G5", "fashion_compact": "G6", "fashion_wide": "G5", }, "baseline_aligned": { "mnist_compact": "G8", "mnist_wide": "G5", "fashion_compact": "G8", "fashion_wide": "G5", }, } # Base Candidate Geometry Assignment (retained for backward compatibility) BASE_GEOMETRIES: Dict[str, str] = GEOMETRY_POLICIES["recovered"] BASELINE_METHODS: Dict[str, str] = { "mnist_compact": "G8", "mnist_wide": "G5", "fashion_compact": "G8", "fashion_wide": "G5", } def compute_latent_dim(total_dim: int, ratio: float) -> int: """ Computes exact model-relative latent dimension from total parameter count and ratio. Validates 0 < ratio <= 1.0 and uses deterministic half-up rounding floor(D*ratio+0.5). Ensures dimension is at least 1 and capped at total_dim. """ if not (0.0 < ratio <= 1.0): raise ValueError(f"ratio must be in (0, 1], got {ratio}") dim = int(math.floor(total_dim * ratio + 0.5)) return max(1, min(total_dim, dim)) def derive_projection_seed( workload_id: str, ratio: float, seed: int, projection_salt: str = "", mode: str = DEFAULT_PROJECTION_SEED_MODE, projection_seed_mode: Optional[str] = None, projection_seed: Optional[Union[int, Dict[str, int]]] = None, ) -> int: """ Derives a deterministic 32-bit projection seed from workload_id, ratio, swarm seed, optional projection_salt, and projection seed mode ('coupled', 'fixed', or 'explicit'). In 'explicit' mode, returns projection_seed directly (or projection_seed[workload_id] if a dict). """ if projection_seed_mode is not None: mode = projection_seed_mode validate_projection_seed_config(mode, projection_seed) if mode == "explicit": if isinstance(projection_seed, dict): if workload_id not in projection_seed: raise ValueError(f"Missing explicit projection_seed for workload '{workload_id}'") return projection_seed[workload_id] assert projection_seed is not None return projection_seed parts = [workload_id, f"{ratio:.5f}"] if mode != "fixed": parts.append(str(seed)) if projection_salt: parts.append(projection_salt) key = ":".join(parts).encode("utf-8") h = hashlib.sha256(key).hexdigest() return int(h[:8], 16) % (2**31 - 1) def construct_equalized_geometry( base_geom: V6GeometryConfig, total_dim: int, latent_dim: int, projection_seed: int, ratio_str: str = "", geometry_multiplier: float = DEFAULT_GEOMETRY_MULTIPLIER, ) -> V6GeometryConfig: """ Constructs an equalized subspace geometry configuration from a base geometry (G6 or G5). Multiplies position radius, launch velocity radius, mutation reset radius, and reflective bound by sqrt(total_dim / latent_dim) * geometry_multiplier to hold decoded per-parameter variance constant. """ validate_geometry_multiplier(geometry_multiplier) scale_factor = math.sqrt(total_dim / latent_dim) if latent_dim < total_dim else 1.0 effective_mult = scale_factor * float(geometry_multiplier) return V6GeometryConfig( config_id=f"{base_geom.config_id}_eq_{ratio_str}", scale_type=base_geom.scale_type, init_position_mode=base_geom.init_position_mode, position_radius=float(base_geom.position_radius * effective_mult), initial_velocity_radius=float(base_geom.initial_velocity_radius * effective_mult), mutation_prob=base_geom.mutation_prob, reset_velocity_radius=float(base_geom.reset_velocity_radius * effective_mult), reflective_bound=float(base_geom.reflective_bound * effective_mult), projection_seed=projection_seed, latent_dim=latent_dim, description=f"Equalized subspace (ratio={ratio_str}, d={latent_dim}/{total_dim})", ) def compute_core_swarm_state_bytes(particles: int, latent_dim: int) -> int: """ Computes exact core swarm state memory footprint in bytes. 5 tensors (Z, V, M, V_sq, P) of size (particles, latent_dim) float32 (4 bytes/elem). """ return 5 * particles * latent_dim * 4 def compute_baseline_core_swarm_state_bytes( workload_id: str, particles: int, total_dim: int, ) -> int: """Match the retained baseline method's persistent core-state accounting.""" particle_states = 5 * particles if BASELINE_METHODS[workload_id] == "G8": particle_states += 1 return particle_states * total_dim * 4 def run_heavy_pso_autoresearch( ratios: Sequence[float] = DEFAULT_RATIOS, particles: int = DEFAULT_PARTICLES, epochs: int = DEFAULT_EPOCHS, subset_size: int = DEFAULT_SUBSET_SIZE, seeds: Sequence[int] = DEFAULT_SEEDS, geometry_policy: str = DEFAULT_GEOMETRY_POLICY, device_str: Optional[str] = None, cache_dir: Optional[Path] = None, output_path: Optional[Path] = None, split_seed: int = DEFAULT_SPLIT_SEED, projection_salt: str = "", projection_scope: Union[str, Dict[str, str]] = DEFAULT_PROJECTION_SCOPE, projection_seed_mode: str = DEFAULT_PROJECTION_SEED_MODE, projection_seed: Optional[Union[int, Dict[str, int]]] = None, geometry_multiplier: float = DEFAULT_GEOMETRY_MULTIPLIER, ) -> Dict[str, Any]: """ Executes the heavy task autoresearch experiment candidate runs across candidate ratios, workloads, and seeds. Aggregates metrics and atomically writes the candidate JSON artifact. """ validate_geometry_multiplier(geometry_multiplier) validate_projection_seed_config(projection_seed_mode, projection_seed) validate_projection_scope_config(projection_scope) if geometry_policy not in GEOMETRY_POLICIES: raise ValueError( f"Invalid geometry_policy '{geometry_policy}'. Must be one of {list(GEOMETRY_POLICIES.keys())}" ) policy_geometries = GEOMETRY_POLICIES[geometry_policy] start_time = time.time() device = resolve_execution_device(device_str) hardware_info = get_hardware_provenance(device) geom_table = get_v6_geometry_table() if cache_dir is None: cache_dir = REPO_ROOT / "result" / "cache" data_cache: Dict[str, Any] = {} workload_meta: Dict[str, Any] = {} # Pre-load datasets and models metadata for wl_id, wl_cfg in WORKLOADS.items(): if wl_cfg.dataset_name not in data_cache: x_search, y_search, x_val, y_val, nested_subsets, data_fp, provenance = prepare_heavy_task_data( dataset_name=wl_cfg.dataset_name, split_seed=split_seed, cache_dir=cache_dir, ) data_cache[wl_cfg.dataset_name] = { "x_search": x_search, "y_search": y_search, "x_val": x_val, "y_val": y_val, "nested_subsets": nested_subsets, "data_fp": data_fp, "provenance": provenance, } bm = create_model(wl_cfg.model_name, seed=41) param_count = sum(p.numel() for p in bm.parameters()) model_fp = compute_model_fingerprint(bm) d_info = data_cache[wl_cfg.dataset_name] eff_proj_seed = ( projection_seed[wl_id] if isinstance(projection_seed, dict) else projection_seed ) workload_meta[wl_id] = { "workload_id": wl_id, "dataset_name": wl_cfg.dataset_name, "model_name": wl_cfg.model_name, "parameter_count": param_count, "total_dim": param_count, "geometry_policy": geometry_policy, "geometry_multiplier": float(geometry_multiplier), "projection_scope": get_effective_projection_scope(projection_scope, wl_id), "projection_seed_mode": str(projection_seed_mode), "projection_seed": eff_proj_seed, "base_geometry_id": policy_geometries[wl_id], "model_fingerprint": model_fp, "data_fingerprint": d_info["data_fp"], "split_fingerprint": d_info["provenance"]["split_fingerprint"], "description": wl_cfg.description, } candidate_runs: Dict[str, Dict[str, Any]] = {} total_runs_executed = 0 total_queries_executed = 0 total_samples_evaluated = 0 for r in ratios: ratio_id = format_ratio_id( r, geometry_policy=geometry_policy, projection_scope=projection_scope, projection_seed_mode=projection_seed_mode, projection_seed=projection_seed, geometry_multiplier=geometry_multiplier, ) wl_candidates: Dict[str, Any] = {} for wl_id, wl_cfg in WORKLOADS.items(): d_info = data_cache[wl_cfg.dataset_name] x_search, y_search = d_info["x_search"], d_info["y_search"] x_val, y_val = d_info["x_val"], d_info["y_val"] nested_subsets = d_info["nested_subsets"] base_model = create_model(wl_cfg.model_name, seed=41) total_dim = sum(p.numel() for p in base_model.parameters()) latent_dim = compute_latent_dim(total_dim, r) base_geom_id = policy_geometries[wl_id] base_geom = geom_table[base_geom_id] state_bytes = compute_core_swarm_state_bytes(particles, latent_dim) baseline_state_bytes = compute_baseline_core_swarm_state_bytes( workload_id=wl_id, particles=particles, total_dim=total_dim, ) per_seed_runs: List[Dict[str, Any]] = [] for s in seeds: proj_seed = derive_projection_seed( wl_id, r, s, projection_salt=projection_salt, mode=projection_seed_mode, projection_seed=projection_seed, ) geom_cfg = construct_equalized_geometry( base_geom=base_geom, total_dim=total_dim, latent_dim=latent_dim, projection_seed=proj_seed, ratio_str=ratio_id, geometry_multiplier=geometry_multiplier, ) eff_scope = get_effective_projection_scope(projection_scope, wl_id) if eff_scope == "tensor_local": transform = TensorLocalLatentTransform(base_model, geom_cfg, device) elif eff_scope == "balanced_global": transform = BalancedGlobalLatentTransform(base_model, geom_cfg, device) elif eff_scope == "two_hash_global": transform = TwoHashGlobalLatentTransform(base_model, geom_cfg, device) elif eff_scope == "largest_tensor_hash": transform = LargestTensorHashLatentTransform(base_model, geom_cfg, device) elif eff_scope == "largest_tensor_row_hash": transform = LargestTensorRowHashLatentTransform(base_model, geom_cfg, device) elif eff_scope == "adjacent_pair": transform = AdjacentPairLatentTransform(base_model, geom_cfg, device) elif eff_scope == "adjacent_difference": transform = AdjacentDifferenceLatentTransform(base_model, geom_cfg, device) elif eff_scope == "global": transform = V6LatentTransform(base_model, geom_cfg, device) else: raise ValueError(f"Invalid effective projection_scope '{eff_scope}' for workload '{wl_id}'") res = run_v6_pso( transform=transform, base_model=base_model, x_search=x_search, y_search=y_search, x_val=x_val, y_val=y_val, nested_subsets=nested_subsets, schedule_str=f"{subset_size}:{epochs}", epochs=epochs, swarm_size=particles, seed=s, device=device, geom_config=geom_cfg, val_check_interval=10, ) opt_time = max(float(res["optimization_wall_time_sec"]), 1e-6) sps = round(float(res["total_sample_evaluations"]) / opt_time, 2) val_metrics_raw = artifact_safe_run(res.get("val_metrics", {})) has_valid_val_metrics = ( isinstance(val_metrics_raw, dict) and len(val_metrics_raw) > 0 and all( isinstance(v, (int, float)) and math.isfinite(float(v)) for v in val_metrics_raw.values() ) ) scalar_metrics = [ res.get("val_selected_loss"), res.get("val_selected_acc"), res.get("gbest_loss"), res.get("gbest_acc"), res.get("wall_time_sec"), res.get("optimization_wall_time_sec"), res.get("validation_wall_time_sec"), sps, ] has_valid_scalars = all( v is not None and isinstance(v, (int, float)) and math.isfinite(float(v)) for v in scalar_metrics ) is_finite = bool(has_valid_val_metrics and has_valid_scalars) seed_record = { "seed": int(s), "projection_seed": int(proj_seed), "projection_salt": str(projection_salt), "projection_scope": get_effective_projection_scope(projection_scope, wl_id), "projection_seed_mode": str(projection_seed_mode), "geometry_multiplier": float(geometry_multiplier), "val_selected_loss": float(res["val_selected_loss"]), "val_selected_acc": float(res["val_selected_acc"]), "val_metrics": val_metrics_raw, "gbest_loss": float(res["gbest_loss"]), "gbest_acc": float(res["gbest_acc"]), "wall_time_sec": float(res["wall_time_sec"]), "optimization_wall_time_sec": float(res["optimization_wall_time_sec"]), "validation_wall_time_sec": float(res["validation_wall_time_sec"]), "total_queries": int(res["total_queries"]), "total_sample_evaluations": int(res["total_sample_evaluations"]), "validation_evaluations": int(res["validation_evaluations"]), "official_test_evaluations": 0, "core_swarm_state_bytes": int(state_bytes), "throughput_samples_per_sec": float(sps), "is_finite": is_finite, } per_seed_runs.append(seed_record) total_runs_executed += 1 total_queries_executed += int(res["total_queries"]) total_samples_evaluated += int(res["total_sample_evaluations"]) # Statistics calculation across seeds acc_list = [r_entry["val_selected_acc"] for r_entry in per_seed_runs] nll_list = [r_entry["val_selected_loss"] for r_entry in per_seed_runs] brier_list = [r_entry["val_metrics"]["brier"] for r_entry in per_seed_runs if r_entry.get("val_metrics")] ece_list = [r_entry["val_metrics"]["ece"] for r_entry in per_seed_runs if r_entry.get("val_metrics")] g_loss_list = [r_entry["gbest_loss"] for r_entry in per_seed_runs] g_acc_list = [r_entry["gbest_acc"] for r_entry in per_seed_runs] wall_list = [r_entry["wall_time_sec"] for r_entry in per_seed_runs] opt_wall_list = [r_entry["optimization_wall_time_sec"] for r_entry in per_seed_runs] sps_list = [r_entry["throughput_samples_per_sec"] for r_entry in per_seed_runs] stats = { "val_acc": calc_stats(acc_list), "val_nll": calc_stats(nll_list), "val_brier": calc_stats(brier_list) if brier_list else None, "val_ece": calc_stats(ece_list) if ece_list else None, "gbest_loss": calc_stats(g_loss_list), "gbest_acc": calc_stats(g_acc_list), "wall_time_sec": calc_stats(wall_list), "optimization_wall_time_sec": calc_stats(opt_wall_list), "throughput_samples_per_sec": calc_stats(sps_list), } state_ratio = float(state_bytes / baseline_state_bytes) wl_candidates[wl_id] = { "candidate_id": ratio_id, "ratio": float(r), "workload_id": wl_id, "geometry_policy": geometry_policy, "geometry_multiplier": float(geometry_multiplier), "projection_scope": get_effective_projection_scope(projection_scope, wl_id), "projection_seed_mode": str(projection_seed_mode), "projection_seed": ( projection_seed[wl_id] if isinstance(projection_seed, dict) else projection_seed ), "base_geometry_id": base_geom_id, "total_dim": total_dim, "latent_dim": latent_dim, "state_ratio": state_ratio, "subset_size": subset_size, "particles": particles, "epochs": epochs, "seeds": list(seeds), "core_swarm_state_bytes": state_bytes, "baseline_core_swarm_state_bytes": baseline_state_bytes, "split_seed": int(split_seed), "data_fingerprint": d_info["data_fp"], "split_fingerprint": d_info["provenance"]["split_fingerprint"], "stats": stats, "per_seed_runs": per_seed_runs, } candidate_runs[ratio_id] = wl_candidates total_wall_time = round(time.time() - start_time, 4) payload = { "protocol_version": AUTORESEARCH_PROTOCOL_VERSION, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "hardware": hardware_info, "official_test_data_loaded": False, "official_test_evaluations": 0, "experiment_config": { "geometry_policy": geometry_policy, "geometry_multiplier": float(geometry_multiplier), "projection_scope": projection_scope, "projection_seed_mode": str(projection_seed_mode), "projection_seed": projection_seed, "projection_salt": str(projection_salt), "ratios": [float(r) for r in ratios], "particles": int(particles), "epochs": int(epochs), "subset_size": int(subset_size), "seeds": [int(s) for s in seeds], "split_seed": int(split_seed), "total_runs": total_runs_executed, "total_queries": total_queries_executed, "total_sample_evaluations": total_samples_evaluated, "total_wall_time_sec": total_wall_time, }, "workloads": workload_meta, "candidate_runs": candidate_runs, } if output_path is not None: save_json_atomic(payload, Path(output_path)) return payload def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="MNIST / FashionMNIST PSO Heavy Autoresearch Candidate Experiment Runner" ) parser.add_argument("--device", type=str, default=None, help="Device (cpu, mps, cuda)") parser.add_argument("--cache-dir", type=str, default=None, help="Dataset cache directory") parser.add_argument( "--split-seed", type=int, default=DEFAULT_SPLIT_SEED, help=f"Dataset train/validation split seed (default: {DEFAULT_SPLIT_SEED})", ) parser.add_argument( "--output", type=str, default="benchmark_results/pso_v6_heavy_autoresearch_candidates.json", help="Output candidate JSON path", ) parser.add_argument( "--geometry-policy", type=str, default=DEFAULT_GEOMETRY_POLICY, choices=list(GEOMETRY_POLICIES.keys()), help="Geometry policy ('recovered' or 'baseline_aligned')", ) parser.add_argument( "--ratios", type=str, default="1,0.5,0.25,0.125,0.03125", help="Comma-separated latent subspace ratios", ) parser.add_argument("--particles", type=int, default=DEFAULT_PARTICLES, help="Swarm size") parser.add_argument("--epochs", type=int, default=DEFAULT_EPOCHS, help="PSO epochs") parser.add_argument( "--subset-size", type=int, default=DEFAULT_SUBSET_SIZE, help="Training subset size", ) parser.add_argument("--seeds", type=str, default="101,102,103", help="Comma-separated seeds") parser.add_argument( "--projection-salt", type=str, default="", help="Optional salt string for projection seed derivation (default: empty)", ) parser.add_argument( "--projection-scope", type=parse_projection_scope_arg, default=DEFAULT_PROJECTION_SCOPE, help="Projection scope ('global', 'tensor_local', 'balanced_global', 'two_hash_global', 'largest_tensor_hash', 'largest_tensor_row_hash', 'adjacent_pair', 'adjacent_difference', or workload dict)", ) parser.add_argument( "--projection-seed-mode", type=str, default=DEFAULT_PROJECTION_SEED_MODE, choices=list(PROJECTION_SEED_MODES), help="Projection seed mode ('coupled', 'fixed', or 'explicit')", ) parser.add_argument( "--projection-seed", type=parse_projection_seed_arg, default=None, help="Exact nonnegative projection seed (int or dict) for explicit mode", ) parser.add_argument( "--geometry-multiplier", type=float, default=DEFAULT_GEOMETRY_MULTIPLIER, help="Geometry radius/bound multiplier (default: 1.0)", ) return parser def main(): parser = build_parser() args = parser.parse_args() ratios = [float(x.strip()) for x in args.ratios.split(",") if x.strip()] seeds = [int(x.strip()) for x in args.seeds.split(",") if x.strip()] cache_path = Path(args.cache_dir) if args.cache_dir else None out_path = Path(args.output) if args.output else None run_heavy_pso_autoresearch( ratios=ratios, split_seed=args.split_seed, projection_seed_mode=args.projection_seed_mode, projection_seed=args.projection_seed, particles=args.particles, epochs=args.epochs, subset_size=args.subset_size, seeds=seeds, geometry_policy=args.geometry_policy, device_str=args.device, cache_dir=cache_path, output_path=out_path, projection_salt=args.projection_salt, projection_scope=args.projection_scope, geometry_multiplier=args.geometry_multiplier, ) if __name__ == "__main__": main()