model.py

from __future__ import annotations

import torch
from torch import nn


class TopKSparseAutoencoder(nn.Module):
    def __init__(
        self, input_dim: int = 96, features: int = 384, top_k: int = 16
    ) -> None:
        super().__init__()
        self.input_dim = input_dim
        self.features = features
        self.top_k = top_k
        self.encoder = nn.Linear(input_dim, features)
        self.decoder = nn.Linear(features, input_dim, bias=False)
        nn.init.normal_(self.decoder.weight, std=0.05)
        self.normalize_dictionary()

    @torch.no_grad()
    def normalize_dictionary(self) -> None:
        norms = self.decoder.weight.norm(dim=0, keepdim=True).clamp_min(1e-8)
        self.decoder.weight.div_(norms)

    def encode(self, activations: torch.Tensor) -> torch.Tensor:
        preactivations = self.encoder(activations)
        values, indices = torch.topk(
            preactivations, self.top_k, dim=1
        )
        values = torch.relu(values)
        sparse = torch.zeros_like(preactivations)
        return sparse.scatter(1, indices, values)

    def forward(
        self, activations: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        features = self.encode(activations)
        return self.decoder(features), features


def parameter_count(module: nn.Module) -> int:
    return sum(parameter.numel() for parameter in module.parameters())