1 hybrid extraction framework (#2)

* first version

* fix mypy & minor adjust

* setup basic ci

Author: kamil271e

.github/actions/setup/action.yml

@@ -0,0 +1,7 @@
+runs:
+  using: composite
+  steps:
+    - name: "Install UV"
+      shell: bash
+      run: |
+        curl -LsSf https://astral.sh/uv/install.sh | sh

.github/workflows/code-quality.yml

@@ -0,0 +1,47 @@
+name: Python Code Quality
+
+on:
+    pull_request:
+    push:
+      branches:
+        - main
+
+jobs:
+  lock_file:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+      - run: uv lock --locked
+
+  linting:
+    runs-on: ubuntu-latest
+    needs: [lock_file]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+      - run: uvx ruff check src/
+
+  formatting:
+    runs-on: ubuntu-latest
+    needs: [lock_file]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+      - run: uvx ruff format --check src/
+
+  type_consistency:
+    runs-on: ubuntu-latest
+    needs: [lock_file]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+      - run: uv run mypy src/
+
+  build:
+    runs-on: [ubuntu-latest]
+    needs: [linting, formatting, type_consistency]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/setup
+      - run: uv build

Makefile

@@ -2,5 +2,8 @@ format:
 	uv run ruff format
 	uv run ruff check --fix src/
 
+check_format:
+	uv run ruff check src
+	
 mypy:
-	uv run mypy src/ --namespace-packages --explicit-package-bases
+	uv run mypy src/

src/autoencoder.py

@@ -0,0 +1,38 @@
+from typing import Any
+
+import torch
+from torch import nn
+
+layers = [1024, 512, 256, 128, 256, 512, 1024]
+
+
+class Autoencoder(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+        encoder_layers: list[nn.Module] = []
+        for i in range(len(layers) // 2):
+            encoder_layers.append(nn.Linear(layers[i], layers[i + 1]))
+            encoder_layers.append(nn.Sigmoid())
+        self.encoder = nn.Sequential(*encoder_layers)
+
+        decoder_layers: list[nn.Module] = []
+        for i in range(len(layers) // 2, len(layers) - 1):
+            decoder_layers.append(nn.Linear(layers[i], layers[i + 1]))
+            decoder_layers.append(nn.Sigmoid())
+        self.decoder = nn.Sequential(*decoder_layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor | Any:
+        x = self.encoder(x)
+        return self.decoder(x)
+
+    def print_weight_shapes(self) -> None:
+        print("Encoder layers:")
+        for layer in self.encoder:
+            if isinstance(layer, nn.Linear):
+                print(layer.weight.shape)
+
+        print("Decoder layers:")
+        for layer in self.decoder:
+            if isinstance(layer, nn.Linear):
+                print(layer.weight.shape)

src/classic_extractor.py

@@ -0,0 +1,8 @@
+import torch
+
+
+class ClassicExtractor:
+    def __init__(self) -> None: ...
+
+    def extract(self, x: list[torch.Tensor]) -> list[torch.Tensor]:
+        return x

src/extraction_framework.py

@@ -0,0 +1,129 @@
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from torch.linalg import svd
+
+from src.autoencoder import layers
+from src.classic_extractor import ClassicExtractor
+from src.inverse_activation import stable_logit
+from src.svd_classifier import SVDClassifier
+from src.utils import restore_weights_from_svd
+
+
+class ExtractionFramework:
+    def __init__(
+        self,
+        decomposition_clf: SVDClassifier,
+        classic_extractor: ClassicExtractor,
+        q_matrix: torch.Tensor,
+    ) -> None:
+        self.decomposition_clf = decomposition_clf
+        self.classic_extractor = classic_extractor
+        self.q_matrix = q_matrix * 0.1  # Scale initial Q-matrix
+
+        # Perform initial decomposition and prediction
+        self.u: torch.Tensor
+        self.s: torch.Tensor
+        self.v: torch.Tensor
+        self.u, self.s, self.v = self.decompose()
+        self.h = self.decomposition_clf.predict(self.s)
+
+        # Initialize placeholders for extracted weights
+        self.w_svd: list[torch.Tensor] = []
+        self.w_hat: list[torch.Tensor] = []
+
+    def decompose(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | Any:
+        """Decomposes the Q-matrix using SVD."""
+        return svd(self.q_matrix)
+
+    def inverse_activation(
+        self, x: torch.Tensor, fun: Callable[..., torch.Tensor] = stable_logit
+    ) -> torch.Tensor:
+        # src/extraction_framework.py:42: error: Too many arguments  [call-arg]
+        """Inverse activation function. Inverse sigmoid by default."""
+        self._clamp_eps = 1e-7
+        self._max_abs = 20.0
+        return fun(x, self._clamp_eps, self._max_abs)
+
+    def extract(self) -> list[torch.Tensor]:
+        """Extracts the weights from the neural network."""
+        h = 0  # Initialize hidden layer dimension tracker
+        w_hat_svd = []  # List to store extracted weights via SVD
+
+        while h < self.h:
+            self.h = h if h != 0 else self.h
+            print(f"Hidden size: {self.h}")
+
+            self.u, self.s, self.v = self.decompose()
+            w_hat_i = restore_weights_from_svd(self.u, self.s, self.h)
+            print(f"W_hat_{self.h} shape: {w_hat_i.shape}")
+            w_hat_svd.append(w_hat_i)
+
+            # Update Q-matrix for next iteration
+            print(f"Q_matrix shape before update: {self.q_matrix.shape}")
+
+            # Ensure data types match for matrix multiplication
+            w_hat_i = w_hat_i.to(self.q_matrix.dtype)
+            try:
+                self.q_matrix = torch.matmul(w_hat_i.transpose(0, 1), self.q_matrix)
+                self.q_matrix = self.inverse_activation(self.q_matrix)
+
+                # Add stability check
+                if torch.isnan(self.q_matrix).any() or torch.isinf(self.q_matrix).any():
+                    self.q_matrix = torch.nan_to_num(
+                        self.q_matrix,
+                        nan=0.0,
+                        posinf=self._max_abs,
+                        neginf=-self._max_abs,
+                    )
+
+            except RuntimeError as e:
+                print(f"Error during Q-matrix update: {e}")
+                # Add recovery mechanism
+                self.q_matrix = torch.clamp(
+                    self.q_matrix, -self._max_abs, self._max_abs
+                )
+                break
+
+            print(f"Q_matrix shape after update: {self.q_matrix.shape}")
+            # Predict next hidden layer dimension
+            h = self.decomposition_clf.predict(self.s)
+
+        w_hat_svd.reverse()
+
+        if not hasattr(self, "w_remaining"):
+            self.initialize_remaining_layers()
+
+        extracted_layers = self.classic_extractor.extract(self.w_remaining)
+
+        print(f"Size of extracted layers: {len(extracted_layers)}")
+        for _ in range(len(extracted_layers)):
+            print(f"Extracted layer {_ + 1} shape: {extracted_layers[_].shape}")
+
+        self.w_hat = extracted_layers + w_hat_svd
+
+        return self.w_hat
+
+    def initialize_remaining_layers(self) -> None:
+        """Initializes remaining layers with random weights."""
+        try:
+            remaining_layers_sizes = [
+                (layers[i + 1], layers[i]) for i in range((len(layers) - 1) // 2)
+            ]
+            self.w_remaining = [
+                torch.randn(size).float() for size in remaining_layers_sizes
+            ]
+        except Exception as e:  # noqa: BLE001
+            print(f"Error initializing remaining layers: {e}")
+
+    def _validate_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Ensure numerical validity of matrix."""
+        tensor = torch.nan_to_num(tensor, nan=0.0)
+        return torch.clamp(tensor, -self._max_abs, self._max_abs)
+
+    def _safe_matmul(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        """Stable matrix multiplication with pre-check."""
+        a = self._validate_tensor(a)
+        b = self._validate_tensor(b)
+        return torch.matmul(a, b)

src/inverse_activation.py

@@ -0,0 +1,13 @@
+import torch
+
+
+def stable_logit(x: torch.Tensor, eps: float, max_abs: float) -> torch.Tensor:
+    """Numerically stable inverse sigmoid with gradient control."""
+    # Clamp probabilities to safe range
+    x_clamped = torch.clamp(x, eps, 1 - eps)
+
+    # Compute logit with separate log operations
+    logit = torch.log(x_clamped) - torch.log1p(-x_clamped)  # More stable than x/(1-x)
+
+    # Prevent gradient explosion
+    return torch.clamp(logit, max_abs, max_abs)

src/main.py

@@ -0,0 +1,52 @@
+from src.autoencoder import Autoencoder
+from src.classic_extractor import ClassicExtractor
+from src.extraction_framework import ExtractionFramework
+from src.svd_classifier import SVDClassifier
+from src.utils import generate_q_matrix, scale_tensor
+
+if __name__ == "__main__":
+    clf = SVDClassifier()
+    extractor = ClassicExtractor()
+    model = Autoencoder()
+
+    Q = generate_q_matrix(model, 1024, 2000)
+    Q = scale_tensor(Q)
+
+    framework = ExtractionFramework(clf, extractor, Q)
+
+    W = framework.extract()
+    # Hidden size: 512
+    # W_hat_512 shape: torch.Size([1024, 512])
+    # Q_matrix shape before update: torch.Size([1024, 2000])
+    # Q_matrix shape after update: torch.Size([512, 2000])
+    # Hidden size: 256
+    # W_hat_256 shape: torch.Size([512, 256])
+    # Q_matrix shape before update: torch.Size([512, 2000])
+    # Q_matrix shape after update: torch.Size([256, 2000])
+    # Hidden size: 128
+    # W_hat_128 shape: torch.Size([256, 128])
+    # Q_matrix shape before update: torch.Size([256, 2000])
+    # Q_matrix shape after update: torch.Size([128, 2000])
+    # Size of extracted layers: 3
+    # Extracted layer 1 shape: torch.Size([512, 1024])
+    # Extracted layer 2 shape: torch.Size([256, 512])
+    # Extracted layer 3 shape: torch.Size([128, 256])
+
+    for w in W:
+        print(w.size())
+    # torch.Size([512, 1024])
+    # torch.Size([256, 512])
+    # torch.Size([128, 256])
+    # torch.Size([256, 128])
+    # torch.Size([512, 256])
+    # torch.Size([1024, 512])
+
+    model.print_weight_shapes()
+    # Encoder layers:
+    # torch.Size([512, 1024])
+    # torch.Size([256, 512])
+    # torch.Size([128, 256])
+    # Decoder layers:
+    # torch.Size([256, 128])
+    # torch.Size([512, 256])
+    # torch.Size([1024, 512])

src/softmax_bottleneck_model.py

@@ -63,3 +63,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor | Any:
         x = self.fc1(x)
         x = self.softmax(x)
         return self.fc2(x)
+
+
+# TODO: use only torch tensors # noqa: FIX002

src/svd_classifier.py

@@ -0,0 +1,12 @@
+import torch
+
+from src.autoencoder import layers
+
+
+class SVDClassifier:
+    def __init__(self) -> None:
+        self.state = 1
+
+    def predict(self, sigma: torch.Tensor) -> int:  # noqa: ARG002
+        self.state += 1
+        return layers[self.state - 1]  # it always return ground truth value