tmnt 0.7.59__py3-none-any.whl → 0.7.61__py3-none-any.whl

tmnt/__init__.py

@@ -1,11 +1,9 @@
 # coding: utf-8
 
-import os
-from .distribution import *
+
 from .preprocess import *
-#from .models import *
+from .sparse import *
 from .utils import *
+from .distribution import *
 
-os.environ["MXNET_STORAGE_FALLBACK_LOG_VERBOSE"] = "0"
-
-__all__ = distribution.__all__ + preprocess.__all__ + utils.__all__
+__all__ = distribution.__all__ + preprocess.__all__ + utils.__all__ # + sparse.__all__

tmnt/estimator.py

@@ -15,7 +15,7 @@ import numpy as np
 import scipy.sparse as sp
 import json
 
-from sklearn.metrics import average_precision_score, top_k_accuracy_score, roc_auc_score, ndcg_score, precision_recall_fscore_support
+from sklearn.metrics import average_precision_score, top_k_accuracy_score, roc_auc_score, ndcg_score
 from tmnt.data_loading import PairedDataLoader, SingletonWrapperLoader, SparseDataLoader, get_llm_model
 from tmnt.modeling import BowVAEModel, SeqBowVED, BaseVAE
 from tmnt.modeling import CrossBatchCosineSimilarityLoss, GeneralizedSDMLLoss, MultiNegativeCrossEntropyLoss, MetricSeqBowVED, MetricBowVAEModel
@@ -29,18 +29,15 @@ from torcheval.metrics import MultilabelAUPRC, MulticlassAUPRC
 ## huggingface specifics
 from transformers.trainer_pt_utils import get_parameter_names
 from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
-from transformers.optimization import AdamW, get_scheduler
+from transformers.optimization import get_scheduler
 
 ## model selection
 import optuna
 
-from itertools import cycle
 import pickle
 from typing import List, Tuple, Dict, Optional, Union, NoReturn
 
 import torch
-from torch.utils.data import Dataset, DataLoader
-from tqdm import tqdm
 
 MAX_DESIGN_MATRIX = 250000000
 

tmnt/sparse/__init__.py

@@ -0,0 +1,12 @@
+# coding: utf-8
+"""
+Copyright (c) 2019 The MITRE Corporation.
+"""
+
+
+from .config import *
+from .estimator import *
+from .inference import *
+from .modeling import *
+
+__all__ = config.__all__ + estimator.__all__ + inference.__all__ + modeling.__all__

tmnt/sparse/config.py

@@ -0,0 +1,35 @@
+import torch 
+
+__all__ = ['get_default_cfg']
+
+def get_default_cfg():
+    default_cfg = {
+        "seed": 49,
+        "batch_size": 4096,
+        "lr": 3e-4,
+        "num_samples": int(1e9),
+        "l1_coeff": 0,
+        "beta1": 0.9,
+        "beta2": 0.99,
+        "max_grad_norm": 100000,
+        "seq_len": 128,
+        "dtype": torch.float32,
+        "site": "resid_pre",
+        "layer": 8,
+        "act_size": 768,
+        "dict_size": 12288,
+        "device": "cuda:0",
+        "input_unit_norm": True,
+        "perf_log_freq": 1000,
+        "sae_type": "topk",
+        "checkpoint_freq": 10000,
+        "n_batches_to_dead": 5,
+
+        # (Batch)TopKSAE specific
+        "top_k": 32,
+        "top_k_aux": 512,
+        "aux_penalty": (1/32),
+        # for jumprelu
+        "bandwidth": 0.001,
+    }
+    return default_cfg

tmnt/sparse/estimator.py

@@ -0,0 +1,98 @@
+import torch
+from torch.utils.data import DataLoader
+from datasets import Dataset, IterableDataset
+import tqdm
+from datasets.arrow_writer import ArrowWriter
+from tmnt.inference import SeqVEDInferencer
+import io, json
+from tmnt.sparse.modeling import BaseAutoencoder
+from typing import List
+
+__all__ = ['ActivationsStore', 'build_activation_store', 'build_activation_store_batching', 'train_sparse_encoder_decoder']
+
+class ActivationsStore:
+    def __init__(
+        self,
+        cfg: dict,
+    ):
+        self.device = cfg["device"]
+        self.activation_path = cfg["activation_path"]
+        shuffle = cfg.get("shuffle_data", False)
+        #self.dataset = Dataset.from_file(self.activation_path).with_format('torch', device=self.device)
+        self.dataset = Dataset.from_file(self.activation_path).select_columns(['data']).shuffle(seed=42).with_format('torch', device=self.device)
+        self.dataloader = DataLoader(self.dataset, 
+                batch_size=cfg["batch_size"], shuffle=shuffle)
+        self.dataloader_iter = iter(self.dataloader)
+        self.cfg = cfg
+
+    def next_batch(self):
+        try:
+            return next(self.dataloader_iter)['data']
+        except (StopIteration, AttributeError):
+            self.dataloader_iter = iter(self.dataloader)
+            return next(self.dataloader_iter)['data']
+
+def build_activation_store(json_input_texts, emb_model_path, arrow_output, max_seq_len=512, json_txt_key='text', device='cpu'):
+
+    inferencer = SeqVEDInferencer.from_saved(emb_model_path, max_length=max_seq_len, device=device)
+    with io.open(json_input_texts) as fp:
+        with ArrowWriter(path=arrow_output) as writer:
+            for l in fp:
+                js = json.loads(l)
+                tokenization_result = inferencer.prep_text(js[json_txt_key])
+                llm_out = inferencer.model.llm(tokenization_result['input_ids'].to(inferencer.device), 
+                                            tokenization_result['attention_mask'].to(inferencer.device))
+                cls_vec = inferencer.model._get_embedding(llm_out, tokenization_result['attention_mask'].to(inferencer.device))
+                enc : List[float] = cls_vec.cpu().detach()[0].tolist() 
+                writer.write({'data': enc})
+            writer.finalize()
+
+def build_activation_store_batching(json_input_texts, emb_model_path, arrow_output, max_seq_len=512, batch_size=42, json_txt_key='text', device='cpu'):
+    inferencer = SeqVEDInferencer.from_saved(emb_model_path, max_length=max_seq_len, device=device)
+    def encode_batch(txt_batch):
+        tokenization_result = inferencer.prep_text(txt_batch)
+        llm_out = inferencer.model.llm(tokenization_result['input_ids'].to(inferencer.device), 
+                                            tokenization_result['attention_mask'].to(inferencer.device))
+        cls_vec = inferencer.model._get_embedding(llm_out, tokenization_result['attention_mask'].to(inferencer.device))
+        encs : List[List[float]] = cls_vec.cpu().detach().tolist() 
+        return zip(txt_batch, encs)
+    
+    def write_encodings(writer: ArrowWriter, txt_enc_pairs):
+        for (t, e) in txt_enc_pairs:
+            writer.write({'text': t, 'data': e})
+
+    with io.open(json_input_texts) as fp:
+        with ArrowWriter(path=arrow_output) as writer:
+            txt_batch = []
+            for l in fp:
+                js = json.loads(l)
+                txt_batch.append(js[json_txt_key])
+                if len(txt_batch) >= batch_size:
+                    encodings = encode_batch(txt_batch)
+                    write_encodings(writer, encodings)
+                    txt_batch = []
+            if len(txt_batch) > 0:
+                encodings = encode_batch(txt_batch)
+                write_encodings(writer, encodings) 
+            writer.finalize() 
+
+
+def train_sparse_encoder_decoder(sed: BaseAutoencoder, activation_store: ActivationsStore, cfg: dict):
+    num_batches = cfg["num_samples"] // cfg["batch_size"]
+    optimizer = torch.optim.Adam(sed.parameters(), lr=cfg["lr"], betas=(cfg["beta1"], cfg["beta2"]))
+    pbar = tqdm.trange(num_batches)
+
+    for i in pbar:
+        batch = activation_store.next_batch()
+        sed_output = sed(batch)
+
+        loss = sed_output["loss"]
+        pbar.set_postfix({"Loss": f"{loss.item():.4f}", "Dead": f"{sed_output['num_dead_features']:.4f}", "L0": f"{sed_output['l0_norm']:.4f}", "L2": f"{sed_output['l2_loss']:.4f}", "L1": f"{sed_output['l1_loss']:.4f}", "L1_norm": f"{sed_output['l1_norm']:.4f}"})
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(sed.parameters(), cfg["max_grad_norm"])
+        sed.make_decoder_weights_and_grad_unit_norm()
+        optimizer.step()
+        optimizer.zero_grad()
+
+
+    

tmnt/sparse/inference.py

@@ -0,0 +1,55 @@
+from tmnt.inference import SeqVEDInferencer
+from scipy.sparse import csr_matrix
+import numpy as np
+from typing import List, Tuple
+from tmnt.distribution import ConceptLogisticGaussianDistribution
+import torch
+from torch.utils.data import DataLoader
+from datasets import Dataset, IterableDataset
+import tqdm
+from datasets.arrow_writer import ArrowWriter
+from tmnt.inference import SeqVEDInferencer
+import io, json
+
+__all__ = ['batch_process_to_arrow']
+
+def csr_to_indices_data(csr_mat):
+    return [ (csr_mat.getrow(ri).indices, csr_mat.getrow(ri).data) for ri in range(csr_mat.shape[0]) ]
+
+def batch_process_to_arrow(model_path, json_input_texts, output_db_path, max_seq_len=512, device='cuda', batch_size=200, json_txt_key='text'):
+
+    inferencer = SeqVEDInferencer.from_saved(model_path, max_length=max_seq_len, device=device)
+    def encode_batch(txt_batch):
+        tokenization_result = inferencer.prep_text(txt_batch)
+        llm_out = inferencer.model.llm(tokenization_result['input_ids'].to(inferencer.device), 
+                                            tokenization_result['attention_mask'].to(inferencer.device))
+        cls_vecs = inferencer.model._get_embedding(llm_out, tokenization_result['attention_mask'].to(inferencer.device))
+        raw_concepts = inferencer.model.latent_distribution.get_sparse_encoding(cls_vecs).cpu().detach()
+        mu_emb = inferencer.model.latent_distribution.get_mu_encoding(cls_vecs)
+        encs : List[List[float]] = cls_vecs.cpu().detach().tolist() 
+        sparse_concepts : List[Tuple[List[int], List[float]]] = csr_to_indices_data(csr_matrix(raw_concepts))
+        topic_embeddings : List[List[float]] = mu_emb.cpu().detach().tolist()
+        print("Lengths: {}, {}, {}, {}".format(len(txt_batch), len(encs), len(sparse_concepts), len(topic_embeddings)))
+        return zip(txt_batch, encs, sparse_concepts, topic_embeddings)
+    
+    def write_encodings(writer: ArrowWriter, txt_enc_pairs):
+        for (text, embedding, sparse_indices_and_data, topic_embedding) in txt_enc_pairs:
+            writer.write({'text': text, 'embedding': embedding, 'indices': sparse_indices_and_data[0], 
+                          'values': sparse_indices_and_data[1], 'topic_embedding': topic_embedding})
+
+    with io.open(json_input_texts) as fp:
+        with ArrowWriter(path=output_db_path) as writer:
+            txt_batch = []
+            for l in fp:
+                js = json.loads(l)
+                txt_batch.append(js[json_txt_key])
+                if len(txt_batch) >= batch_size:
+                    encodings = encode_batch(txt_batch)
+                    write_encodings(writer, encodings)
+                    txt_batch = []
+            if len(txt_batch) > 0:
+                encodings = encode_batch(txt_batch)
+                write_encodings(writer, encodings) 
+            writer.finalize() 
+    
+

tmnt/sparse/modeling.py

@@ -0,0 +1,385 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.autograd as autograd
+
+
+__all__ = ['BatchTopKEncoder', 'BatchTopKSAE', 'TopKEncoder', 'TopKSAE', 'VanillaEncoder', 'VanillaSAE', 'JumpReLUEncoder', 'JumpReLUSAE']
+
+class BaseEncoder(nn.Module):
+    
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        torch.manual_seed(self.cfg['seed'])
+        self.b_enc = nn.Parameter(torch.zeros(self.cfg['dict_size']))
+        self.W_enc = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(self.cfg['act_size'], self.cfg['dict_size'])
+            )
+        )
+        self.to(cfg['dtype']).to(cfg['device'])
+
+    def get_dict_size(self):
+        return int(self.cfg['dict_size'])
+
+    def preprocess_input(self, x):
+        if self.cfg.get("input_unit_norm", False):
+            x_mean = x.mean(dim=-1, keepdim=True)
+            x = x - x_mean
+            x_std = x.std(dim=-1, keepdim=True)
+            x = x / (x_std + 1e-5)
+            return x, x_mean, x_std
+        else:
+            return x, None, None
+        
+
+class BaseAutoencoder(nn.Module):
+    """Base class for autoencoder models."""
+
+    def __init__(self, cfg, encoder: BaseEncoder):
+        super().__init__()
+
+        self.cfg = cfg
+        torch.manual_seed(self.cfg["seed"])
+        
+        self.encoder = encoder
+
+        self.b_dec = nn.Parameter(torch.zeros(self.cfg["act_size"]))
+        self.W_dec = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(self.cfg["dict_size"], self.cfg["act_size"])
+            )
+        )
+        self.W_dec.data[:] = self.encoder.W_enc.t().data
+        self.W_dec.data[:] = self.W_dec / self.W_dec.norm(dim=-1, keepdim=True)
+        self.num_batches_not_active = torch.zeros((self.cfg["dict_size"],)).to(
+            cfg["device"]
+        )
+
+        self.to(cfg["dtype"]).to(cfg["device"])
+
+    def postprocess_output(self, x_reconstruct, x_mean, x_std):
+        if self.cfg.get("input_unit_norm", False):
+            x_reconstruct = x_reconstruct * x_std + x_mean
+        return x_reconstruct
+
+    @torch.no_grad()
+    def make_decoder_weights_and_grad_unit_norm(self):
+        W_dec_normed = self.W_dec / self.W_dec.norm(dim=-1, keepdim=True)
+        W_dec_grad_proj = (self.W_dec.grad * W_dec_normed).sum(
+            -1, keepdim=True
+        ) * W_dec_normed
+        self.W_dec.grad -= W_dec_grad_proj
+        self.W_dec.data = W_dec_normed
+
+    def update_inactive_features(self, acts):
+        self.num_batches_not_active += (acts.sum(0) == 0).float()
+        self.num_batches_not_active[acts.sum(0) > 0] = 0
+
+
+class BatchTopKEncoder(BaseEncoder):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def forward(self, x):
+        x, x_mean, x_std = self.preprocess_input(x)
+
+        acts = F.relu(x @ self.W_enc)
+        acts_topk = torch.topk(acts.flatten(), self.cfg["top_k"] * x.shape[0], dim=-1)
+        acts_topk = (
+            torch.zeros_like(acts.flatten())
+            .scatter(-1, acts_topk.indices, acts_topk.values)
+            .reshape(acts.shape)
+        )
+        return acts, acts_topk, x, x_mean, x_std
+
+
+
+class BatchTopKSAE(BaseAutoencoder):
+    def __init__(self, cfg:dict , encoder: BatchTopKEncoder):
+        super().__init__(cfg, encoder)
+
+    def forward(self, x):
+        acts, acts_topk, x, x_mean, x_std = self.encoder(x)
+        x_reconstruct = acts_topk @ self.W_dec + self.b_dec
+
+        self.update_inactive_features(acts_topk)
+        output = self.get_loss_dict(x, x_reconstruct, acts, acts_topk, x_mean, x_std)
+        return output
+
+
+    def get_loss_dict(self, x, x_reconstruct, acts, acts_topk, x_mean, x_std):
+        l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
+        l1_norm = acts_topk.float().abs().sum(-1).mean()
+        l1_loss = self.cfg["l1_coeff"] * l1_norm
+        l0_norm = (acts_topk > 0).float().sum(-1).mean()
+        aux_loss = self.get_auxiliary_loss(x, x_reconstruct, acts)
+        loss = l2_loss + l1_loss + aux_loss
+        num_dead_features = (
+            self.num_batches_not_active > self.cfg["n_batches_to_dead"]
+        ).sum()
+        sae_out = self.postprocess_output(x_reconstruct, x_mean, x_std)
+        output = {
+            "sae_out": sae_out,
+            "feature_acts": acts_topk,
+            "num_dead_features": num_dead_features,
+            "loss": loss,
+            "l1_loss": l1_loss,
+            "l2_loss": l2_loss,
+            "l0_norm": l0_norm,
+            "l1_norm": l1_norm,
+            "aux_loss": aux_loss,
+        }
+        return output
+
+    def get_auxiliary_loss(self, x, x_reconstruct, acts):
+        dead_features = self.num_batches_not_active >= self.cfg["n_batches_to_dead"]
+        if dead_features.sum() > 0:
+            residual = x.float() - x_reconstruct.float()
+            acts_topk_aux = torch.topk(
+                acts[:, dead_features],
+                min(self.cfg["top_k_aux"], dead_features.sum()),
+                dim=-1,
+            )
+            acts_aux = torch.zeros_like(acts[:, dead_features]).scatter(
+                -1, acts_topk_aux.indices, acts_topk_aux.values
+            )
+            x_reconstruct_aux = acts_aux @ self.W_dec[dead_features]
+            l2_loss_aux = (
+                self.cfg["aux_penalty"]
+                * (x_reconstruct_aux.float() - residual.float()).pow(2).mean()
+            )
+            return l2_loss_aux
+        else:
+            return torch.tensor(0, dtype=x.dtype, device=x.device)
+
+
+class TopKEncoder(BaseEncoder):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def forward(self, x):
+        x, x_mean, x_std = self.preprocess_input(x)
+        acts = F.relu(x @ self.W_enc)
+        acts_topk = torch.topk(acts, self.cfg["top_k"], dim=-1)
+        acts_topk = torch.zeros_like(acts).scatter(
+            -1, acts_topk.indices, acts_topk.values
+        )
+        return acts, acts_topk, x, x_mean, x_std 
+
+class TopKSAE(BaseAutoencoder):
+    def __init__(self, cfg: dict, encoder: TopKEncoder):
+        super().__init__(cfg, encoder)
+
+    def forward(self, x):
+        acts, acts_topk, x, x_mean, x_std = self.encoder(x)
+        x_reconstruct = acts_topk @ self.W_dec + self.b_dec
+        self.update_inactive_features(acts_topk)
+        output = self.get_loss_dict(x, x_reconstruct, acts, acts_topk, x_mean, x_std)
+        return output
+
+    def get_loss_dict(self, x, x_reconstruct, acts, acts_topk, x_mean, x_std):
+        l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
+        l1_norm = acts_topk.float().abs().sum(-1).mean()
+        l1_loss = self.cfg["l1_coeff"] * l1_norm
+        l0_norm = (acts_topk > 0).float().sum(-1).mean()
+        aux_loss = self.get_auxiliary_loss(x, x_reconstruct, acts)
+        loss = l2_loss + l1_loss + aux_loss
+        num_dead_features = (
+            self.num_batches_not_active > self.cfg["n_batches_to_dead"]
+        ).sum()
+        sae_out = self.postprocess_output(x_reconstruct, x_mean, x_std)
+        output = {
+            "sae_out": sae_out,
+            "feature_acts": acts_topk,
+            "num_dead_features": num_dead_features,
+            "loss": loss,
+            "l1_loss": l1_loss,
+            "l2_loss": l2_loss,
+            "l0_norm": l0_norm,
+            "l1_norm": l1_norm,
+            "aux_loss": aux_loss,
+        }
+        return output
+
+    def get_auxiliary_loss(self, x, x_reconstruct, acts):
+        dead_features = self.num_batches_not_active >= self.cfg["n_batches_to_dead"]
+        if dead_features.sum() > 0:
+            residual = x.float() - x_reconstruct.float()
+            acts_topk_aux = torch.topk(
+                acts[:, dead_features],
+                min(self.cfg["top_k_aux"], dead_features.sum()),
+                dim=-1,
+            )
+            acts_aux = torch.zeros_like(acts[:, dead_features]).scatter(
+                -1, acts_topk_aux.indices, acts_topk_aux.values
+            )
+            x_reconstruct_aux = acts_aux @ self.W_dec[dead_features]
+            l2_loss_aux = (
+                self.cfg["aux_penalty"]
+                * (x_reconstruct_aux.float() - residual.float()).pow(2).mean()
+            )
+            return l2_loss_aux
+        else:
+            return torch.tensor(0, dtype=x.dtype, device=x.device)
+        
+
+class VanillaEncoder(BaseEncoder):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def forward(self, x): 
+        x, x_mean, x_std = self.preprocess_input(x)
+        acts = F.relu(x @ self.W_enc + self.b_enc)
+        return acts, x, x_mean, x_std
+
+
+class VanillaSAE(BaseAutoencoder):
+    def __init__(self, cfg, encoder: VanillaEncoder):
+        super().__init__(cfg, encoder)
+
+    def forward(self, x):
+        acts, x, x_mean, x_std = self.encoder(x)
+        x_reconstruct = acts @ self.W_dec + self.b_dec
+        self.update_inactive_features(acts)
+        output = self.get_loss_dict(x, x_reconstruct, acts, x_mean, x_std)
+        return output
+
+    def get_loss_dict(self, x, x_reconstruct, acts, x_mean, x_std):
+        l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
+        l1_norm = acts.float().abs().sum(-1).mean()
+        l1_loss = self.cfg["l1_coeff"] * l1_norm
+        l0_norm = (acts > 0).float().sum(-1).mean()
+        loss = l2_loss + l1_loss
+        num_dead_features = (
+            self.num_batches_not_active > self.cfg["n_batches_to_dead"]
+        ).sum()
+
+        sae_out = self.postprocess_output(x_reconstruct, x_mean, x_std)
+        output = {
+            "sae_out": sae_out,
+            "feature_acts": acts,
+            "num_dead_features": num_dead_features,
+            "loss": loss,
+            "l1_loss": l1_loss,
+            "l2_loss": l2_loss,
+            "l0_norm": l0_norm,
+            "l1_norm": l1_norm,
+        }
+        return output
+
+import torch
+import torch.nn as nn
+import torch.autograd as autograd
+
+class RectangleFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return ((x > -0.5) & (x < 0.5)).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (x,) = ctx.saved_tensors
+        grad_input = grad_output.clone()
+        grad_input[(x <= -0.5) | (x >= 0.5)] = 0
+        return grad_input
+
+class JumpReLUFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x, log_threshold, bandwidth):
+        ctx.save_for_backward(x, log_threshold, torch.tensor(bandwidth))
+        threshold = torch.exp(log_threshold)
+        return x * (x > threshold).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, log_threshold, bandwidth_tensor = ctx.saved_tensors
+        bandwidth = bandwidth_tensor.item()
+        threshold = torch.exp(log_threshold)
+        x_grad = (x > threshold).float() * grad_output
+        threshold_grad = (
+            -(threshold / bandwidth)
+            * RectangleFunction.apply((x - threshold) / bandwidth)
+            * grad_output
+        )
+        return x_grad, threshold_grad, None  # None for bandwidth
+
+class JumpReLU(nn.Module):
+    def __init__(self, feature_size, bandwidth, device='cpu'):
+        super(JumpReLU, self).__init__()
+        self.log_threshold = nn.Parameter(torch.zeros(feature_size, device=device))
+        self.bandwidth = bandwidth
+
+    def forward(self, x):
+        return JumpReLUFunction.apply(x, self.log_threshold, self.bandwidth)
+
+class StepFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x, log_threshold, bandwidth):
+        ctx.save_for_backward(x, log_threshold, torch.tensor(bandwidth))
+        threshold = torch.exp(log_threshold)
+        return (x > threshold).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, log_threshold, bandwidth_tensor = ctx.saved_tensors
+        bandwidth = bandwidth_tensor.item()
+        threshold = torch.exp(log_threshold)
+        x_grad = torch.zeros_like(x)
+        threshold_grad = (
+            -(1.0 / bandwidth)
+            * RectangleFunction.apply((x - threshold) / bandwidth)
+            * grad_output
+        )
+        return x_grad, threshold_grad, None  # None for bandwidth
+    
+
+class JumpReLUEncoder(BaseEncoder):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+
+    def forward(self, x):
+        x, x_mean, x_std = self.preprocess_input(x)
+
+        pre_activations = torch.relu(x @ self.W_enc + self.b_enc)
+        feature_magnitudes = self.jumprelu(pre_activations)
+        return feature_magnitudes, x, x_mean, x_std
+
+
+class JumpReLUSAE(BaseAutoencoder):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+        self.jumprelu = JumpReLU(feature_size=cfg["dict_size"], bandwidth=cfg["bandwidth"], device=cfg["device"])
+
+    def forward(self, x):
+        feature_magnitudes, x, x_mean, x_std = self.encoder(x)
+        x_reconstructed = feature_magnitudes @ self.W_dec + self.b_dec
+
+        return self.get_loss_dict(x, x_reconstructed, feature_magnitudes, x_mean, x_std)
+
+    def get_loss_dict(self, x, x_reconstruct, acts, x_mean, x_std):
+        l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
+
+        l0 = StepFunction.apply(acts, self.jumprelu.log_threshold, self.cfg["bandwidth"]).sum(dim=-1).mean()
+        l0_loss = self.cfg["l1_coeff"] * l0
+        l1_loss = l0_loss
+
+        loss = l2_loss + l1_loss
+        num_dead_features = (
+            self.num_batches_not_active > self.cfg["n_batches_to_dead"]
+        ).sum()
+
+        sae_out = self.postprocess_output(x_reconstruct, x_mean, x_std)
+        output = {
+            "sae_out": sae_out,
+            "feature_acts": acts,
+            "num_dead_features": num_dead_features,
+            "loss": loss,
+            "l1_loss": l1_loss,
+            "l2_loss": l2_loss,
+            "l0_norm": l0,
+            "l1_norm": l0,
+        }
+        return output

{tmnt-0.7.59.dist-info → tmnt-0.7.61.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tmnt
-Version: 0.7.59
+Version: 0.7.61
 Summary: Topic modeling neural toolkit
 Home-page: https://github.com/mitre/tmnt.git
 Author: The MITRE Corporation
@@ -48,7 +48,7 @@ Dynamic: summary
 The Topic Modeling Neural Toolkit (TMNT) is a software library that enables training
 topic models as neural network-based variational auto-encoders.
 
-Current stable version is: 0.7.54
+Current stable version is: 0.7.61
 
 Documentation can be found here: https://tmnt.readthedocs.io/en/stable/
 

{tmnt-0.7.59.dist-info → tmnt-0.7.61.dist-info}/RECORD

@@ -1,14 +1,19 @@
-tmnt/__init__.py,sha256=EPNq1H7UMyMewWT_zTGBaC7ZouvCywX_gMX4G1dtmvw,250
+tmnt/__init__.py,sha256=3W6sRKacvIMQuxGZpscjU52B4sDcasEA1RCXJVHaK-Q,203
 tmnt/configuration.py,sha256=P8PEhzVPKO5xG0FrdTLRQ60OYWigbzPY-OSx_hzQlrY,10054
 tmnt/data_loading.py,sha256=LcVcXX00UsuAillRPILcvmqj3AsCIgzB6V_S6lfsbIY,19335
 tmnt/distribution.py,sha256=4gn1wnszVAErzICCvZXSYki0G78WC3_jyBr27N-Aj3E,15108
-tmnt/estimator.py,sha256=KnnvSNXm6cRL0GwDrGdgqqPX5ZubpCQ0WqcSXJDkUU4,68072
+tmnt/estimator.py,sha256=SJW9koDHb_lTwmqXm4Ilgga7MMMmRGLd3ylEFYu3bZg,67933
 tmnt/eval_npmi.py,sha256=8S-IE-bEhtQofF6oKeXs7oaUeu-7yDlaEqjMj52gmNQ,6549
 tmnt/inference.py,sha256=Iwc2_w7QrS1epiVEm_Ewx5sYFNNMDfvhMJETOgJqm0E,15783
 tmnt/modeling.py,sha256=rGHQsW7ldycFUd1f9NzcnNuSRElr600vLwmYPl6YY0M,30215
 tmnt/preprocess/__init__.py,sha256=gwMejkQrnqKS05i0JVsUru2hDUR5jE1hKC10dL934GU,170
 tmnt/preprocess/tokenizer.py,sha256=-ZgowfbHrM040vbNTktZM_hdl6HDTqxSJ4mDAxq3dUs,14050
 tmnt/preprocess/vectorizer.py,sha256=RaianZ_DG3Nc-RI96FtmI4PCZPi5Nipx9a5xndLZ52M,20689
+tmnt/sparse/__init__.py,sha256=BEhOm_o0UrVUKTG3rSiBJzE7qQQL9HRSZ1MHCA2GJu8,249
+tmnt/sparse/config.py,sha256=DzqzmclxbUmmBZHS7poMicsvgOZCMXTXWEY5bTpPrRI,827
+tmnt/sparse/estimator.py,sha256=PSqIopCZG7WO5MXYcN-CAs7R7M36jZXjQeUp0I5kG9Q,4721
+tmnt/sparse/inference.py,sha256=KiSUjo_lO14qkUzzo9HuOWZ5S5oJRzoSsSa2dOg54UY,2850
+tmnt/sparse/modeling.py,sha256=hy2nkI2tq3Bpmj7OfXr_zwHRajWBwsqz7T3wHTuEbyE,13753
 tmnt/utils/__init__.py,sha256=1PZsxRPsHI_DnOpxD0iAhLxhxHnx6Svzg3W-79YfWWs,237
 tmnt/utils/csv2json.py,sha256=A1TXy-uxA4dc9tw0tjiHzL7fv4C6b0Uc_bwI1keTmKU,795
 tmnt/utils/log_utils.py,sha256=ZtR4nF_Iee23ev935YQcTtXv-cCC7lgXkXLl_yokfS4,2075
@@ -18,9 +23,9 @@ tmnt/utils/pubmed_utils.py,sha256=3sHwoun7vxb0GV-arhpXLMUbAZne0huAh9xQNy6H40E,12
 tmnt/utils/random.py,sha256=qY75WG3peWoMh9pUyCPBEo6q8IvkF6VRjeb5CqJOBF8,327
 tmnt/utils/recalibrate.py,sha256=TmpB8An8bslICZ13UTJfIvr8VoqiSedtpHxec4n8CHk,1439
 tmnt/utils/vocab.py,sha256=J6GFGLyvDgdmtVQjYlyzWjuykRD3kllCKPG1z0lI0P8,3504
-tmnt-0.7.59.dist-info/licenses/LICENSE,sha256=qFZJrfJ7Zi4IXDiyiGVrHWic_l1h2tc36tI8Z7rK9bs,11356
-tmnt-0.7.59.dist-info/licenses/NOTICE,sha256=p0kYIVAkReTFaGb4C-qPa7h5ztze6hGzOpjCMMbOipU,425
-tmnt-0.7.59.dist-info/METADATA,sha256=WnXkGjITudOxSCMTfqQVMdx2UvfxN4O-SLvdYVtTckY,1663
-tmnt-0.7.59.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-tmnt-0.7.59.dist-info/top_level.txt,sha256=RpYgUl187sXnqmiwKjZZdcDlHz2AALs6bGdUcukyd_E,5
-tmnt-0.7.59.dist-info/RECORD,,
+tmnt-0.7.61.dist-info/licenses/LICENSE,sha256=qFZJrfJ7Zi4IXDiyiGVrHWic_l1h2tc36tI8Z7rK9bs,11356
+tmnt-0.7.61.dist-info/licenses/NOTICE,sha256=p0kYIVAkReTFaGb4C-qPa7h5ztze6hGzOpjCMMbOipU,425
+tmnt-0.7.61.dist-info/METADATA,sha256=hyptC3YXersP-7oR89NgKU4UksTrvsbcC5Ml38WSYG0,1663
+tmnt-0.7.61.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+tmnt-0.7.61.dist-info/top_level.txt,sha256=RpYgUl187sXnqmiwKjZZdcDlHz2AALs6bGdUcukyd_E,5
+tmnt-0.7.61.dist-info/RECORD,,

{tmnt-0.7.59.dist-info → tmnt-0.7.61.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any