| | """ |
| | Copy of the existing SubspaceFeaturizer implementation for submission. |
| | This file provides the same SubspaceFeaturizer functionality in a self-contained format. |
| | """ |
| |
|
| | import torch |
| | import torch.nn as nn |
| | import pyvene as pv |
| | from CausalAbstraction.model_units.model_units import Featurizer |
| |
|
| |
|
| | class SubspaceFeaturizerModuleCopy(torch.nn.Module): |
| | def __init__(self, rotate_layer): |
| | super().__init__() |
| | self.rotate = rotate_layer |
| | |
| | def forward(self, x): |
| | r = self.rotate.weight.T |
| | f = x.to(r.dtype) @ r.T |
| | error = x - (f @ r).to(x.dtype) |
| | return f, error |
| |
|
| |
|
| | class SubspaceInverseFeaturizerModuleCopy(torch.nn.Module): |
| | def __init__(self, rotate_layer): |
| | super().__init__() |
| | self.rotate = rotate_layer |
| | |
| | def forward(self, f, error): |
| | r = self.rotate.weight.T |
| | return (f.to(r.dtype) @ r).to(f.dtype) + error.to(f.dtype) |
| |
|
| |
|
| | class SubspaceFeaturizerCopy(Featurizer): |
| | def __init__(self, shape=None, rotation_subspace=None, trainable=True, id="subspace"): |
| | assert shape is not None or rotation_subspace is not None, "Either shape or rotation_subspace must be provided." |
| | if shape is not None: |
| | self.rotate = pv.models.layers.LowRankRotateLayer(*shape, init_orth=True) |
| | elif rotation_subspace is not None: |
| | shape = rotation_subspace.shape |
| | self.rotate = pv.models.layers.LowRankRotateLayer(*shape, init_orth=False) |
| | self.rotate.weight.data.copy_(rotation_subspace) |
| | self.rotate = torch.nn.utils.parametrizations.orthogonal(self.rotate) |
| |
|
| | if not trainable: |
| | self.rotate.requires_grad_(False) |
| |
|
| | |
| | featurizer = SubspaceFeaturizerModuleCopy(self.rotate) |
| | inverse_featurizer = SubspaceInverseFeaturizerModuleCopy(self.rotate) |
| | |
| | super().__init__(featurizer, inverse_featurizer, n_features=self.rotate.weight.shape[1], id=id) |
