1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 182,
5  "end_line": 232,
6  "content": "class Compressor:\n    \"\"\"Adaptive TopK-EF with layer-aware budget and stochastic exploration.\n\n    Novel contributions over standard TopK-EF:\n    1. Layer-adaptive budget: higher keep ratio for smaller tensors that are\n       disproportionately sensitive to compression information loss.\n    2. Skip tiny tensors: BN/bias params are sent uncompressed (negligible\n       communication cost, significant accuracy preservation).\n    3. Stochastic exploration: multiplicative noise on magnitude scores breaks\n       selection ties and rotates which borderline elements get communicated,\n       improving gradient coverage and reducing dead zones.\n    \"\"\"\n\n    def __init__(self, compress_ratio=0.01):\n        self.compress_ratio = compress_ratio\n        self.residuals = {}\n        self.noise_scale = 0.03  # Exploration noise magnitude\n\n    def compress(self, tensor, name):\n        shape = tensor.shape\n        tensor_flat = tensor.flatten()\n        numel = tensor_flat.numel()\n        # Skip compression for tiny tensors (BN params, biases)\n        if numel < 256:\n            return [tensor_flat.clone()], (numel, shape, False)\n        # Error feedback: add accumulated residual\n        if name in self.residuals:\n            tensor_flat = tensor_flat + self.residuals[name]\n        # Layer-adaptive K: higher relative budget for smaller layers\n        if numel < 2048:\n            ratio = min(0.5, self.compress_ratio * 10)\n        elif numel < 32768:\n            ratio = min(0.15, self.compress_ratio * 3)\n        else:\n            ratio = self.compress_ratio\n        k = max(1, int(numel * ratio))\n        # Stochastic TopK: noise on scores for exploration\n        abs_vals = tensor_flat.abs()\n        noise = 1.0 + self.noise_scale * torch.randn_like(abs_vals)\n        scores = abs_vals * noise.clamp(min=0.5)\n        _, indices = torch.topk(scores, k, sorted=False)\n        values = tensor_flat[indices]\n        # Update residual (what was not communicated)\n        decompressed_flat = torch.zeros_like(tensor_flat)\n        decompressed_flat.scatter_(0, indices, values)\n        self.residuals[name] = tensor_flat - decompressed_flat\n        return [values, indices], (numel, shape, True)\n\n    def decompress(self, compressed_tensors, ctx):\n        numel, shape, was_compressed = ctx\n        if not was_compressed:\n            return compressed_tensors[0].view(shape)\n        values, indices = compressed_tensors\n        t = torch.zeros(numel, dtype=values.dtype, device=values.device)\n        t.scatter_(0, indices, values)\n        return t.view(shape)\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 182,
5  "end_line": 232,
6  "content": "class Compressor:\n    \"\"\"PowerSGD with error feedback and warm-start projection matrices.\n\n    Approximates gradient matrices via a single randomized power iteration\n    with warm-started right-projection Q. Error feedback accumulates the\n    low-rank approximation residual. Rank per layer adapts to matrix\n    dimensions and target compression ratio. Small / 1-D tensors (e.g.\n    batch-norm, biases) bypass compression to avoid high relative error.\n    \"\"\"\n\n    def __init__(self, compress_ratio=0.01):\n        self.compress_ratio = compress_ratio\n        self.residuals = {}\n        self.q_memory = {}\n\n    def compress(self, tensor, name):\n        shape = tensor.shape\n        numel = tensor.numel()\n        # Skip compression for small or 1D tensors\n        if numel < 256 or len(shape) < 2:\n            return [tensor.clone().flatten()], (shape, False)\n        # Reshape to 2D: (out_features, in_features * spatial_dims)\n        matrix = tensor.view(shape[0], -1)\n        # Error feedback: add accumulated residual\n        if name in self.residuals:\n            matrix = matrix + self.residuals[name]\n        m, n = matrix.shape\n        # Adaptive rank from target compression ratio\n        rank = max(2, int(self.compress_ratio * m * n / (m + n)))\n        rank = min(rank, min(m, n))\n        # Warm-start Q from previous iteration, or random init\n        if name in self.q_memory and self.q_memory[name].shape == (n, rank):\n            q = self.q_memory[name]\n        else:\n            q = torch.randn(n, rank, device=tensor.device, dtype=tensor.dtype)\n        # Single power iteration: P = orth(M @ Q), Q_new = M^T @ P\n        p = matrix @ q\n        p, _ = torch.linalg.qr(p)\n        q_new = matrix.t() @ p\n        # Save Q for warm-start next iteration\n        self.q_memory[name] = q_new.clone()\n        # Error feedback: store residual\n        self.residuals[name] = matrix - p @ q_new.t()\n        return [p, q_new], (shape, True)\n\n    def decompress(self, compressed_tensors, ctx):\n        shape, is_compressed = ctx\n        if not is_compressed:\n            return compressed_tensors[0].view(shape)\n        p, q_new = compressed_tensors\n        return (p @ q_new.t()).view(shape)"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 182,
5  "end_line": 232,
6  "content": "class Compressor:\n    \"\"\"PowerSGD with error feedback, warm-start, and enhanced rank.\n    Low-rank approximation via 2 power iterations for 2D+ gradients.\n    Rank adapts per layer. Error feedback corrects residuals.\n    Small/1D tensors bypass compression.\n    \"\"\"\n\n    def __init__(self, compress_ratio=0.01):\n        self.compress_ratio = compress_ratio\n        self.residuals = {}\n        self.q_memory = {}\n\n    def compress(self, tensor, name):\n        shape = tensor.shape\n        numel = tensor.numel()\n        if numel < 256 or len(shape) < 2:\n            return [tensor.clone().flatten()], (shape, False)\n        matrix = tensor.view(shape[0], -1)\n        if name in self.residuals:\n            matrix = matrix + self.residuals[name]\n        m, n = matrix.shape\n        rank = max(4, int(2.0 * self.compress_ratio * m * n / (m + n)))\n        rank = min(rank, min(m, n))\n        if name in self.q_memory and self.q_memory[name].shape == (n, rank):\n            q = self.q_memory[name]\n        else:\n            q = torch.randn(n, rank, device=tensor.device, dtype=tensor.dtype)\n        # Two power iterations for better subspace estimation\n        p = matrix @ q\n        p, _ = torch.linalg.qr(p)\n        q = matrix.t() @ p\n        q, _ = torch.linalg.qr(q)\n        p = matrix @ q\n        p, _ = torch.linalg.qr(p)\n        q_new = matrix.t() @ p\n        # Save orthogonalized Q for stable warm-start\n        q_ws, _ = torch.linalg.qr(q_new)\n        self.q_memory[name] = q_ws\n        self.residuals[name] = matrix - p @ q_new.t()\n        return [p, q_new], (shape, True)\n\n    def decompress(self, compressed_tensors, ctx):\n        shape, is_compressed = ctx\n        if not is_compressed:\n            return compressed_tensors[0].view(shape)\n        p, q_new = compressed_tensors\n        return (p @ q_new.t()).view(shape)\n\n\n\n"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_compressor.py",
4  "start_line": 203,
5  "end_line": 203,
6  "content": "        rank = max(4, int(3.0 * self.compress_ratio * m * n / (m + n)))"
7}