1{
2  "op": "replace",
3  "filename": "pytorch-examples/optimization_parity/custom_strategy.py",
4  "start_line": 220,
5  "end_line": 255,
6  "content": "def init_model(model: nn.Sequential, config: TaskConfig) -> None:\n    \"\"\"Initialize the fixed two-layer MLP without using the hidden secret.\"\"\"\n    for i, layer in enumerate(model):\n        if isinstance(layer, nn.Linear):\n            if i == 0:  # first linear layer\n                # Kaiming He initialization for ReLU\n                nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')\n            else:  # second linear layer\n                # Xavier uniform initialization\n                nn.init.xavier_uniform_(layer.weight, gain=1.0)\n            # Initialize bias to small positive value to encourage activation\n            nn.init.constant_(layer.bias, 0.01)\n\n\ndef make_dataset(\n    secret: tuple[int, ...],\n    config: TaskConfig,\n    seed: int,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Return a structured dataset that systematically covers secret bit patterns.\"\"\"\n    generator = torch.Generator().manual_seed(seed)\n    K = config.secret_size\n    N = config.n_features\n    # Generate all 2^K patterns of secret bits\n    num_patterns = 1 << K  # 2^K\n    # Create tensor of shape (num_patterns, K) with all binary combinations\n    secret_patterns = torch.zeros((num_patterns, K), dtype=torch.float32)\n    for i in range(num_patterns):\n        for j in range(K):\n            secret_patterns[i, j] = float((i >> j) & 1)\n    \n    # Determine repeats to fill up to desired dataset size\n    target_size = 4096  # tunable, but keep moderate for multi-epoch training\n    repeats = max(1, target_size // num_patterns)\n    actual_size = repeats * num_patterns\n    \n    # Expand secret patterns\n    secret_bits = secret_patterns.repeat(repeats, 1)  # shape (actual_size, K)\n    \n    # Generate random noise for non-secret bits\n    non_secret_mask = torch.ones(N, dtype=torch.bool)\n    non_secret_mask[list(secret)] = False\n    num_non_secret = non_secret_mask.sum().item()\n    \n    # Random binary matrix for non-secret bits\n    noise = torch.randint(\n        low=0, high=2,\n        size=(actual_size, num_non_secret),\n        generator=generator,\n        dtype=torch.float32\n    )\n    \n    # Assemble full input vectors\n    x = torch.zeros((actual_size, N), dtype=torch.float32)\n    # Place secret bits at secret positions\n    for idx, sec_idx in enumerate(secret):\n        x[:, sec_idx] = secret_bits[:, idx]\n    # Place noise at non-secret positions\n    x[:, non_secret_mask] = noise\n    \n    # Compute labels\n    y = parity_labels(x, secret)\n    \n    return x, y\n\n\ndef get_optimizer_config(config: TaskConfig) -> dict[str, float]:\n    \"\"\"Return tuned AdamW hyperparameters for sparse parity.\"\"\"\n    return {\n        \"lr\": 5e-4,          # lower learning rate for stability\n        \"wd\": 0.0,           # no weight decay (helps sparse features)\n        \"beta1\": 0.9,        # standard\n        \"beta2\": 0.999,      # standard\n    }"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-examples/optimization_parity/custom_strategy.py",
4  "start_line": 220,
5  "end_line": 293,
6  "content": "def init_model(model: nn.Sequential, config: TaskConfig) -> None:\n    \"\"\"Initialize the fixed two-layer MLP without using the hidden secret.\"\"\"\n    for layer in model:\n        if isinstance(layer, nn.Linear):\n            gain = nn.init.calculate_gain(\"relu\") if layer is model[0] else 1.0\n            nn.init.xavier_uniform_(layer.weight, gain=gain)\n            # Initialize bias to zero (default)\n            nn.init.zeros_(layer.bias)\n\n\ndef make_dataset(\n    secret: tuple[int, ...],\n    config: TaskConfig,\n    seed: int,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Return a dataset with linearly independent secret patterns plus random noise.\"\"\"\n    generator = torch.Generator().manual_seed(seed)\n    K = config.secret_size\n    N = config.n_features\n    \n    # Create K+2 canonical patterns for secret bits:\n    # 1) All zeros\n    # 2) All ones\n    # 3) Each pattern with exactly one secret bit = 1, others = 0 (K patterns)\n    num_patterns = K + 2\n    secret_patterns = torch.zeros((num_patterns, K), dtype=torch.float32)\n    # all zeros already zero\n    # all ones\n    secret_patterns[1, :] = 1.0\n    # single-bit patterns\n    for i in range(K):\n        secret_patterns[2 + i, i] = 1.0\n    \n    # Determine repeats to fill up to desired dataset size\n    target_size = 2048  # moderate size for multi-epoch training\n    repeats = max(1, target_size // num_patterns)\n    actual_size = repeats * num_patterns\n    \n    # Expand secret patterns\n    secret_bits = secret_patterns.repeat(repeats, 1)  # shape (actual_size, K)\n    \n    # Generate random noise for non-secret bits (different for each example)\n    non_secret_mask = torch.ones(N, dtype=torch.bool)\n    non_secret_mask[list(secret)] = False\n    num_non_secret = non_secret_mask.sum().item()\n    \n    # Random binary matrix for non-secret bits\n    noise = torch.randint(\n        low=0, high=2,\n        size=(actual_size, num_non_secret),\n        generator=generator,\n        dtype=torch.float32\n    )\n    \n    # Assemble full input vectors\n    x = torch.zeros((actual_size, N), dtype=torch.float32)\n    # Place secret bits at secret positions\n    for idx, sec_idx in enumerate(secret):\n        x[:, sec_idx] = secret_bits[:, idx]\n    # Place noise at non-secret positions\n    x[:, non_secret_mask] = noise\n    \n    # Compute labels\n    y = parity_labels(x, secret)\n    \n    return x, y\n\n\ndef get_optimizer_config(config: TaskConfig) -> dict[str, float]:\n    \"\"\"Return tuned AdamW hyperparameters.\"\"\"\n    return {\n        \"lr\": 1e-3,          # standard learning rate\n        \"wd\": 1e-4,          # small weight decay for regularization\n        \"beta1\": 0.9,\n        \"beta2\": 0.999,\n    }"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-examples/optimization_parity/custom_strategy.py",
4  "start_line": 220,
5  "end_line": 255,
6  "content": "def init_model(model: nn.Sequential, config: TaskConfig) -> None:\n    \"\"\"Initialize the fixed two-layer MLP without using the hidden secret.\"\"\"\n    for i, layer in enumerate(model):\n        if isinstance(layer, nn.Linear):\n            if i == 0:  # first layer: use Kaiming uniform for ReLU\n                nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')\n            else:  # second layer: use orthogonal initialization\n                nn.init.orthogonal_(layer.weight, gain=1.0)\n            nn.init.zeros_(layer.bias)\n\n\ndef make_dataset(\n    secret: tuple[int, ...],\n    config: TaskConfig,\n    seed: int,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Return a training dataset that covers all 2^K patterns of secret bits.\"\"\"\n    generator = torch.Generator().manual_seed(seed)\n    K = config.secret_size\n    N = config.n_features\n    \n    # We want to cover all 2^K patterns of secret bits.\n    # For each pattern, we'll generate multiple examples with random non-secret bits.\n    # Target total dataset size: use max_train_examples (12.8M)\n    max_examples = config.max_train_examples\n    patterns = 2 ** K  # 256 for K=8\n    \n    # Examples per pattern\n    examples_per_pattern = max_examples // patterns\n    # Ensure at least 1 example per pattern\n    examples_per_pattern = max(1, examples_per_pattern)\n    \n    # Build dataset\n    x_list = []\n    y_list = []\n    \n    # For each pattern of secret bits (0..patterns-1)\n    for pattern in range(patterns):\n        # Generate binary representation of pattern for secret bits\n        secret_bits = [(pattern >> i) & 1 for i in range(K)]\n        \n        # Generate examples_per_pattern samples\n        for _ in range(examples_per_pattern):\n            # Create random binary vector for all N bits\n            vec = torch.randint(\n                low=0, high=2,\n                size=(N,),\n                generator=generator,\n                dtype=torch.int64\n            ).to(torch.float32)\n            \n            # Set secret bits according to pattern\n            for idx, secret_idx in enumerate(secret):\n                vec[secret_idx] = secret_bits[idx]\n            \n            # Compute parity label (should match pattern parity)\n            parity = sum(secret_bits) % 2\n            \n            x_list.append(vec)\n            y_list.append(float(parity))\n    \n    x = torch.stack(x_list)\n    y = torch.tensor(y_list, dtype=torch.float32)\n    \n    # Shuffle the dataset\n    perm = torch.randperm(x.shape[0], generator=generator)\n    x = x[perm]\n    y = y[perm]\n    \n    # Trim to max_examples if we exceeded due to rounding\n    if x.shape[0] > max_examples:\n        x = x[:max_examples]\n        y = y[:max_examples]\n    \n    return x, y\n\n\ndef get_optimizer_config(config: TaskConfig) -> dict[str, float]:\n    \"\"\"Return AdamW hyperparameters tuned for sparse parity.\"\"\"\n    return {\n        \"lr\": 1e-3,      # Learning rate\n        \"wd\": 0.0,       # No weight decay - avoid regularization that may hurt\n        \"beta1\": 0.9,    # Default\n        \"beta2\": 0.99,   # Slightly lower than 0.999 for more stable updates\n    }"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-examples/optimization_parity/custom_strategy.py",
4  "start_line": 220,
5  "end_line": 255,
6  "content": "def init_model(model: nn.Sequential, config: TaskConfig) -> None:\n    \"\"\"Initialize the fixed two-layer MLP without using the hidden secret.\"\"\"\n    for layer in model:\n        if isinstance(layer, nn.Linear):\n            if layer is model[0]:\n                # Use Kaiming normal initialization for ReLU\n                nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')\n                # Initialize bias to small positive values to encourage initial activations\n                nn.init.constant_(layer.bias, 0.01)\n            else:\n                # Output layer: use smaller weights to avoid saturation\n                nn.init.xavier_uniform_(layer.weight, gain=1.0)\n                nn.init.zeros_(layer.bias)\n\n\ndef make_dataset(\n    secret: tuple[int, ...],\n    config: TaskConfig,\n    seed: int,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Return a training dataset that systematically covers all secret bit combinations.\"\"\"\n    generator = torch.Generator().manual_seed(seed)\n    K = config.secret_size\n    N = config.n_features\n    max_examples = config.max_train_examples\n    \n    # All possible combinations of secret bits (2^K)\n    secret_combos = torch.arange(2**K, dtype=torch.int64)\n    # Expand to binary matrix of shape (2^K, K)\n    secret_bits = ((secret_combos.unsqueeze(1) >> torch.arange(K)) & 1).to(torch.float32)\n    \n    # Determine repetitions to fill up to max_examples\n    repeats = max_examples // (2**K)\n    # Ensure at least 1 repeat, but also add some fully random examples for diversity\n    repeats = max(1, repeats - 1)  # reserve space for random examples\n    num_systematic = repeats * (2**K)\n    num_random = max_examples - num_systematic\n    \n    # Systematic examples: for each secret combination, repeat with random background bits\n    secret_indices = torch.tensor(secret, dtype=torch.long)\n    non_secret_mask = torch.ones(N, dtype=torch.bool)\n    non_secret_mask[secret_indices] = False\n    non_secret_indices = torch.where(non_secret_mask)[0]\n    \n    # Generate random background bits for each systematic example\n    total_systematic = repeats * (2**K)\n    # Background matrix shape (total_systematic, N-K)\n    background = torch.randint(\n        low=0, high=2,\n        size=(total_systematic, len(non_secret_indices)),\n        generator=generator,\n        dtype=torch.float32\n    )\n    \n    # Construct x matrix\n    x_systematic = torch.zeros((total_systematic, N), dtype=torch.float32)\n    # Fill secret bits: repeat each secret combination `repeats` times\n    secret_expanded = secret_bits.repeat(repeats, 1)  # shape (total_systematic, K)\n    x_systematic[:, secret_indices] = secret_expanded\n    # Fill background bits\n    x_systematic[:, non_secret_indices] = background\n    \n    # Compute labels\n    y_systematic = (secret_expanded.sum(dim=1) % 2).to(torch.float32)\n    \n    # Add fully random examples for diversity (optional)\n    if num_random > 0:\n        x_random = torch.randint(\n            low=0, high=2,\n            size=(num_random, N),\n            generator=generator,\n            dtype=torch.float32\n        )\n        y_random = parity_labels(x_random, secret)\n        # Concatenate\n        x = torch.cat([x_systematic, x_random], dim=0)\n        y = torch.cat([y_systematic, y_random], dim=0)\n    else:\n        x = x_systematic\n        y = y_systematic\n    \n    # Shuffle the dataset\n    perm = torch.randperm(x.shape[0], generator=generator)\n    x = x[perm]\n    y = y[perm]\n    \n    return x, y\n\n\ndef get_optimizer_config(config: TaskConfig) -> dict[str, float]:\n    \"\"\"Return AdamW hyperparameters optimized for sparse parity learning.\"\"\"\n    return {\n        \"lr\": 2e-2,           # Higher learning rate for faster convergence\n        \"wd\": 1e-4,           # Low weight decay to avoid excessive regularization\n        \"beta1\": 0.9,\n        \"beta2\": 0.999,\n    }"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-examples/optimization_parity/custom_strategy.py",
4  "start_line": 220,
5  "end_line": 316,
6  "content": "def init_model(model: nn.Sequential, config: TaskConfig) -> None:\n    \"\"\"Initialize the fixed two-layer MLP without using the hidden secret.\"\"\"\n    for layer in model:\n        if isinstance(layer, nn.Linear):\n            if layer is model[0]:\n                # Orthogonal initialization for better gradient propagation\n                nn.init.orthogonal_(layer.weight, gain=nn.init.calculate_gain('relu'))\n                # Small positive bias to avoid dead ReLU\n                nn.init.constant_(layer.bias, 0.01)\n            else:\n                # Output layer: Xavier uniform\n                nn.init.xavier_uniform_(layer.weight, gain=1.0)\n                nn.init.zeros_(layer.bias)\n\n\ndef make_dataset(\n    secret: tuple[int, ...],\n    config: TaskConfig,\n    seed: int,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Return a training dataset that systematically covers all secret bit combinations with limited size.\"\"\"\n    generator = torch.Generator().manual_seed(seed)\n    K = config.secret_size\n    N = config.n_features\n    # Use a moderate dataset size to allow multiple epochs\n    target_size = 10_000  # tunable\n    dataset_size = min(target_size, config.max_train_examples)\n    \n    # All possible combinations of secret bits (2^K)\n    secret_combos = torch.arange(2**K, dtype=torch.int64)\n    secret_bits = ((secret_combos.unsqueeze(1) >> torch.arange(K)) & 1).to(torch.float32)  # (2^K, K)\n    \n    # Determine repetitions to fill dataset\n    repeats = dataset_size // (2**K)\n    if repeats == 0:\n        # If dataset smaller than 2^K, sample a subset of combos\n        indices = torch.randperm(2**K, generator=generator)[:dataset_size]\n        secret_subset = secret_bits[indices]\n        repeats = 1\n        total_systematic = dataset_size\n    else:\n        # Use all combos, each repeated `repeats` times\n        secret_subset = secret_bits.repeat(repeats, 1)  # (repeats*2^K, K)\n        total_systematic = repeats * (2**K)\n    \n    # Ensure we don't exceed dataset_size due to rounding\n    if total_systematic > dataset_size:\n        # trim excess\n        secret_subset = secret_subset[:dataset_size]\n        total_systematic = dataset_size\n    \n    secret_indices = torch.tensor(secret, dtype=torch.long)\n    non_secret_mask = torch.ones(N, dtype=torch.bool)\n    non_secret_mask[secret_indices] = False\n    non_secret_indices = torch.where(non_secret_mask)[0]\n    \n    # Generate random background bits for each systematic example\n    background = torch.randint(\n        low=0, high=2,\n        size=(total_systematic, len(non_secret_indices)),\n        generator=generator,\n        dtype=torch.float32\n    )\n    \n    # Construct x matrix\n    x = torch.zeros((total_systematic, N), dtype=torch.float32)\n    x[:, secret_indices] = secret_subset\n    x[:, non_secret_indices] = background\n    \n    # Compute labels\n    y = (secret_subset.sum(dim=1) % 2).to(torch.float32)\n    \n    # Shuffle\n    perm = torch.randperm(x.shape[0], generator=generator)\n    x = x[perm]\n    y = y[perm]\n    \n    return x, y\n\n\ndef get_optimizer_config(config: TaskConfig) -> dict[str, float]:\n    \"\"\"Return AdamW hyperparameters for sparse parity learning.\"\"\"\n    return {\n        \"lr\": 1e-3,\n        \"wd\": 1e-2,\n        \"beta1\": 0.9,\n        \"beta2\": 0.999,\n    }"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-examples/optimization_parity/custom_strategy.py",
4  "start_line": 220,
5  "end_line": 307,
6  "content": "def init_model(model: nn.Sequential, config: TaskConfig) -> None:\n    \"\"\"Initialize the fixed two-layer MLP without using the hidden secret.\"\"\"\n    for layer in model:\n        if isinstance(layer, nn.Linear):\n            if layer is model[0]:\n                # Kaiming uniform initialization with ReLU gain\n                nn.init.kaiming_uniform_(layer.weight, nonlinearity='relu')\n                # Small positive bias to avoid dead ReLU\n                nn.init.constant_(layer.bias, 0.01)\n            else:\n                # Output layer: small weights to avoid saturation\n                nn.init.xavier_uniform_(layer.weight, gain=0.1)\n                nn.init.zeros_(layer.bias)\n\n\ndef make_dataset(\n    secret: tuple[int, ...],\n    config: TaskConfig,\n    seed: int,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Return a training dataset that emphasizes learning parity of secret bits.\"\"\"\n    generator = torch.Generator().manual_seed(seed)\n    K = config.secret_size\n    N = config.n_features\n    # Use moderate dataset size for multiple epochs\n    target_size = 20_000  # larger than 10k but still manageable\n    dataset_size = min(target_size, config.max_train_examples)\n    \n    # All possible combinations of secret bits (2^K)\n    secret_combos = torch.arange(2**K, dtype=torch.int64)\n    secret_bits = ((secret_combos.unsqueeze(1) >> torch.arange(K)) & 1).to(torch.float32)  # (2^K, K)\n    \n    # Determine repetitions to fill dataset\n    repeats = dataset_size // (2**K)\n    if repeats == 0:\n        # If dataset smaller than 2^K, sample a subset of combos\n        indices = torch.randperm(2**K, generator=generator)[:dataset_size]\n        secret_subset = secret_bits[indices]\n        repeats = 1\n        total_systematic = dataset_size\n    else:\n        # Use all combos, each repeated `repeats` times\n        secret_subset = secret_bits.repeat(repeats, 1)  # (repeats*2^K, K)\n        total_systematic = repeats * (2**K)\n    \n    # Ensure we don't exceed dataset_size due to rounding\n    if total_systematic > dataset_size:\n        secret_subset = secret_subset[:dataset_size]\n        total_systematic = dataset_size\n    \n    secret_indices = torch.tensor(secret, dtype=torch.long)\n    non_secret_mask = torch.ones(N, dtype=torch.bool)\n    non_secret_mask[secret_indices] = False\n    non_secret_indices = torch.where(non_secret_mask)[0]\n    \n    # For each systematic example, decide background: zero or random\n    # 80% zero background, 20% random background\n    background_choices = torch.rand(total_systematic, generator=generator)\n    zero_background_mask = background_choices < 0.8\n    \n    # Generate random background bits for those not zero\n    random_background = torch.randint(\n        low=0, high=2,\n        size=(total_systematic, len(non_secret_indices)),\n        generator=generator,\n        dtype=torch.float32\n    )\n    # Zero background\n    zero_background = torch.zeros((total_systematic, len(non_secret_indices)), dtype=torch.float32)\n    \n    # Choose\n    background = torch.where(zero_background_mask.unsqueeze(1), zero_background, random_background)\n    \n    # Construct x matrix\n    x = torch.zeros((total_systematic, N), dtype=torch.float32)\n    x[:, secret_indices] = secret_subset\n    x[:, non_secret_indices] = background\n    \n    # Compute labels\n    y = (secret_subset.sum(dim=1) % 2).to(torch.float32)\n    \n    # Add a small fraction of fully random examples (10%) to encourage invariance\n    num_random = int(0.1 * dataset_size)\n    if num_random > 0:\n        x_rand = torch.randint(\n            low=0, high=2,\n            size=(num_random, N),\n            generator=generator,\n            dtype=torch.float32\n        )\n        y_rand = parity_labels(x_rand, secret)\n        # Concatenate\n        x = torch.cat([x, x_rand], dim=0)\n        y = torch.cat([y, y_rand], dim=0)\n    \n    # Shuffle\n    perm = torch.randperm(x.shape[0], generator=generator)\n    x = x[perm]\n    y = y[perm]\n    \n    return x, y\n\n\ndef get_optimizer_config(config: TaskConfig) -> dict[str, float]:\n    \"\"\"Return AdamW hyperparameters for robust parity learning.\"\"\"\n    return {\n        \"lr\": 5e-4,           # Lower learning rate for stable learning\n        \"wd\": 1e-3,           # Moderate weight decay\n        \"beta1\": 0.9,\n        \"beta2\": 0.999,\n    }"
7}