Overcomplete

https://kempnerinstitute.github.io/overcomplete/
Admin
Overcomplete is a PyTorch research library designed for dictionary learning and sparse autoencoder...
Tokens:31,580
Snippets:279
Trust Score:9.7
Update:1 month ago
Show doc for...
Tokens:
Raw
### Train a TopKSAE model

Source: https://kempnerinstitute.github.io/overcomplete

Example showing how to initialize a TopKSAE model, prepare a DataLoader, and train the model using the provided train_sae function.

```python
import torch
from torch.utils.data import DataLoader, TensorDataset
from overcomplete.sae import TopKSAE, train_sae

Activations = torch.randn(N, d)
sae = TopKSAE(d, nb_concepts=16_000, top_k=10, device='cuda')

dataloader = DataLoader(TensorDataset(Activations), batch_size=1024)
optimizer = torch.optim.Adam(sae.parameters(), lr=5e-4)

def criterion(x, x_hat, pre_codes, codes, dictionary):
  mse = (x - x_hat).square().mean()
  return mse

logs = train_sae(sae, dataloader, criterion, optimizer,
                 nb_epochs=20, device='cuda')
```

--------------------------------

### Install Overcomplete via pip

Source: https://kempnerinstitute.github.io/overcomplete

Use this command to install the Overcomplete library and its dependencies.

```bash
pip install overcomplete
```

--------------------------------

### Visualize top-concept activations with overcomplete.visualization

Source: https://kempnerinstitute.github.io/overcomplete/visualization

Shows how to import the four visualization functions and apply them to a synthetic batch of 100 images and 10,000 concept heatmaps for concept 3. The example uses random tensors and demonstrates the recommended overlay with 'jet' colormap and alpha 0.35, plus evidence, zoom, and contour variants.

```python
from overcomplete.visualization import (overlay_top_heatmaps,
evidence_top_images, zoom_top_images, contour_top_image)

# lets imagine we have 100 images and 10k concepts maps
# and we want to visualize concept 3
images = torch.randn(100, 3, 256, 256)
heatmaps = torch.randn(100, 14, 14, 10_000)

# heatmap + transparency (recommended)
overlay_top_heatmaps(images, heatmaps, concept_id=3,
                     cmap='jet', alpha=0.35)
# transparency based
evidence_top_images(images, heatmaps, concept_id=3)
# zoom into max activating crops
zoom_top_images(images, heatmaps, concept_id=3, zoom_size=100)
# contour of most important part (boundary)
contour_top_image(images, heatmaps, concept_id=3)
```

--------------------------------

### Basic Usage of BatchTopKSAE

Source: https://kempnerinstitute.github.io/overcomplete/saes/batchtopk

Demonstrates initializing the BatchTopKSAE model and accessing the running threshold during evaluation.

```python
from overcomplete import BatchTopKSAE

# define a Batch Top K SAE with input dimension 768, 10k concepts
# and top_k = 50 (for the entire batch!)
sae = BatchTopKSAE(768, 10_000, top_k=50)

# the threshold momentum is used to estimate
# the final threshold (when in eval)
sae = BatchTopKSAE(768, 10_000, top_k=10, threshold_momentum=0.95)
# ... training sae
sae = sae.eval()
# now top_k is no longer use and instead an
# internal threshold is used
print(sae.running_threshold)
```

--------------------------------

### energy_of_codes(codes, dictionary)

Source: https://kempnerinstitute.github.io/overcomplete/metrics

Computes the energy of codes given a dictionary. For example, with X input sample, Z the codes and D the dictionary: X = ZD, Energy(Z) = || E[Z]D ||^2 and corresponds to the average energy the codes bring to the reconstruction.

```APIDOC
## energy_of_codes(codes, dictionary)

### Description
Compute the energy of codes given a dictionary. For example, with X input sample, Z the codes and D the dictionary: X = ZD, Energy(Z) = || E[Z]D ||^2 and correspond to the average energy the codes bring to the reconstruction.

### Parameters
- **codes** (torch.Tensor) - Required - Codes tensor of shape (batch_size, num_codes).
- **dictionary** (torch.Tensor) - Required - Dictionary tensor of shape (num_codes, dim).

### Return
- **torch.Tensor** - Energy of codes, one per codes dimension.
```

--------------------------------

### TopKSAE.__init__

Source: https://kempnerinstitute.github.io/overcomplete/saes/topk

Initializes a TopKSAE instance with specified input shape, number of concepts, and optional parameters for top-k sparsity, encoder module, dictionary parameters, and device.

```APIDOC
## TopKSAE.__init__

### Description
Initializes a TopKSAE instance.

### Method
Constructor

### Signature
```python
TopKSAE(input_shape, nb_concepts, top_k=None, encoder_module=None, dictionary_params=None, device='cpu')
```

### Parameters
- **input_shape** (int or tuple of int) - Required - Dimensionality of the input data, do not include batch dimensions. It is usually 1d (dim), 2d (seq length, dim) or 3d (dim, height, width).
- **nb_concepts** (int) - Required - Number of components/concepts in the dictionary. The dictionary is overcomplete if the number of concepts > in_dimensions.
- **top_k** (int, optional) - Optional - Number of top activations to keep in the latent representation, by default n_components // 10 (sparsity of 90%).
- **encoder_module** (nn.Module or string, optional) - Optional - Custom encoder module, by default None. If None, a simple Linear + BatchNorm default encoder is used. If string, the name of the registered encoder module.
- **dictionary_params** (dict, optional) - Optional - Parameters that will be passed to the dictionary layer. See DictionaryLayer for more details.
- **device** (str, optional) - Optional - Device to run the model on, by default 'cpu'.

### Example
```python
from overcomplete import TopKSAE

# define a TopK SAE with input dimension 768, 10k concepts
sae = TopKSAE(768, 10_000, top_k=5)

# adjust the encoder module (you can also)
# directly pass your own encoder module
sae = TopKSAE(768, 10_000, top_k=10,
              encoder_module='mlp_bn_1')
```
```

--------------------------------

### BatchTopKSAE.__init__

Source: https://kempnerinstitute.github.io/overcomplete/saes/batchtopk

Initializes a BatchTopKSAE instance with specified input shape, number of concepts, top-k value, and optional parameters for threshold momentum, encoder module, dictionary parameters, and device.

```APIDOC
## BatchTopKSAE.__init__

### Description
Initializes a BatchTopKSAE instance.

### Method
Constructor

### Signature
__init__(self, input_shape, nb_concepts, top_k=None, threshold_momentum=0.9, encoder_module=None, dictionary_params=None, device='cpu')

### Parameters
- **input_shape** (int or tuple of int) - Required - Dimensionality of the input data (excluding the batch dimension).
- **nb_concepts** (int) - Required - Number of latent dimensions (components) of the autoencoder.
- **top_k** (int) - Required - The number of activations to keep (the kth highest activation is used as threshold).
- **threshold_momentum** (float, optional) - Optional - Momentum for the running threshold update (default is 0.9).
- **encoder_module** (nn.Module or str, optional) - Optional - Custom encoder module (or its registered name). If None, a default encoder is used.
- **dictionary_params** (dict, optional) - Optional - Parameters that will be passed to the dictionary layer. See DictionaryLayer for more details.
- **device** (str, optional) - Optional - Device on which to run the model (default is 'cpu').

### Example
```python
from overcomplete import BatchTopKSAE
sae = BatchTopKSAE(768, 10_000, top_k=50)
```
```

--------------------------------

### Basic usage of SemiNMF with multiplicative update solver

Source: https://kempnerinstitute.github.io/overcomplete/optimization/seminmf

Shows how to instantiate a SemiNMF model with 10,000 concepts using the 'mu' solver, fit it to input activations A to obtain codes Z and dictionary D, then encode new data and decode (reconstruct) data from codes.

```python
from overcomplete import SemiNMF

# define a Semi-NMF model with 10k concepts using
# the multiplicative update solver
semi_nmf = SemiNMF(nb_concepts=10_000, solver='mu')

# fit the model to input activations A
Z, D = semi_nmf.fit(A)

# encode new data
Z = semi_nmf.encode(A)
# decode (reconstruct) data from codes
A_hat = semi_nmf.decode(Z)
```

--------------------------------

### SemiNMF.__init__

Source: https://kempnerinstitute.github.io/overcomplete/optimization/seminmf

Initializes the SemiNMF model with the specified number of concepts and solver settings.

```APIDOC
## SemiNMF.__init__

### Description
Initializes the SemiNMF model with the specified number of concepts and solver settings.

### Method
Constructor

### Parameters
- **nb_concepts** (int) - Required - Number of components to learn.
- **solver** (str, optional) - Solver to use, by default 'mu'. Possible values are 'mu' (multiplicative update) and 'pgd' (projected gradient descent), 'pgd' allows for sparsity penalty.
- **device** (str, optional) - Device to use for tensor computations, by default 'cpu'.
- **tol** (float, optional) - Tolerance value for the stopping criterion, by default 1e-4.
- **l1_penalty** (float, optional) - L1 penalty for the sparsity constraint, by default 0.0. Only used with 'pgd' solver.
- **verbose** (bool, optional) - Whether to print the status of the optimization, by default False.

### Returns
None
```

--------------------------------

### RATopKSAE.__init__

Source: https://kempnerinstitute.github.io/overcomplete/saes/archetypal_saes

Initializes a Relaxed Archetypal TopK SAE with the given input shape, number of concepts, archetype points, and sparsity parameters.

```APIDOC
## RATopKSAE.__init__

### Description
Initializes a Relaxed Archetypal TopK SAE.

### Method
Constructor

### Signature
`__init__(self, input_shape, nb_concepts, points, top_k=None, delta=1.0, use_multiplier=True, **kwargs)`

### Parameters
- **input_shape** (int) - Required - Dimensionality of the input data (excluding the batch dimension).
- **nb_concepts** (int) - Required - Number of dictionary atoms (concepts).
- **points** (torch.Tensor) - Required - The data points used to initialize/define the archetypes. Shape should be (num_points, input_shape).
- **top_k** (int) - Optional - Number of top activations to keep in the latent representation. By default, 10% sparsity is used.
- **delta** (float) - Optional - Delta parameter for the archetypal dictionary, by default 1.0.
- **use_multiplier** (bool) - Optional - Whether to use a learnable multiplier that parametrize the ball (e.g. if this parameter is 3 then the dictionary atoms are all on the ball of radius 3). By default True.
- **kwargs** (dict) - Optional - Additional arguments passed to the parent TopKSAE (e.g., encoder_module, device).
```

--------------------------------

### Basic NMF usage with overcomplete

Source: https://kempnerinstitute.github.io/overcomplete/optimization/nmf

Shows how to instantiate an NMF model with 10,000 concepts using the HALS solver, fit it to activation matrix A, and then encode and decode data. The fit method returns both codes Z and dictionary D, while encode and decode operate on new data.

```python
from overcomplete import NMF

# define an NMF model with 10k concepts using the HALS solver
nmf = NMF(nb_concepts=10_000, solver='hals')

# fit the model to input activations A
Z, D = nmf.fit(A)

# encode new data
Z = nmf.encode(A)
# decode (reconstruct) data from codes
A_hat = nmf.decode(Z)
```

--------------------------------

### Initialize and encode with OMPSAE

Source: https://kempnerinstitute.github.io/overcomplete/saes/omp

Demonstrates how to instantiate the OMPSAE model and perform an encoding operation. The encode method returns the residual and the sparse codes.

```python
import torch
from overcomplete.sae import OMPSAE

x = torch.randn(64, 512)
sae = OMPSAE(
    input_shape=512,
    nb_concepts=4_096,
    k=4,            # pursuit steps
    max_iter=15,    # NNLS iterations
    dropout=0.1,    # optional atom dropout
    encoder_module="identity",
    device="cuda"
)

residual, codes = sae.encode(x)
```

--------------------------------

### Basic TopKSAE instantiation

Source: https://kempnerinstitute.github.io/overcomplete/saes/topk

Creates a TopKSAE with input dimension 768, 10,000 concepts, and top_k=5. The default encoder is used.

```python
from overcomplete import TopKSAE

# define a TopK SAE with input dimension 768, 10k concepts
sae = TopKSAE(768, 10_000, top_k=5)
```

--------------------------------

### SAE.__init__

Source: https://kempnerinstitute.github.io/overcomplete/saes/vanilla

Initializes a Vanilla SAE model with specified input shape, number of concepts, encoder module, dictionary parameters, and device.

```APIDOC
## SAE.__init__

### Description
Initializes a Vanilla SAE model with specified input shape, number of concepts, encoder module, dictionary parameters, and device.

### Method
Constructor

### Signature
__init__(self, input_shape, nb_concepts, encoder_module=None, dictionary_params=None, device='cpu')

### Parameters
- **input_shape** (int or tuple of int) - Required - Dimensionality of the input data, do not include batch dimensions. It is usually 1d (dim), 2d (seq length, dim) or 3d (dim, height, width).
- **nb_concepts** (int) - Required - Number of components/concepts in the dictionary. The dictionary is overcomplete if the number of concepts > in_dimensions.
- **encoder_module** (nn.Module or string, optional) - Optional - Custom encoder module, by default None. If None, a simple Linear + BatchNorm default encoder is used. If string, the name of the registered encoder module.
- **dictionary_params** (dict, optional) - Optional - Parameters that will be passed to the dictionary layer. See DictionaryLayer for more details.
- **device** (str, optional) - Optional - Device to run the model on, by default 'cpu'.

### Return
None
```

--------------------------------

### Basic usage of JumpSAE with custom kernel and bandwidth

Source: https://kempnerinstitute.github.io/overcomplete/saes/jumprelu

Import the JumpSAE class and create an instance with input dimension 768 and 10,000 concepts. Adjust the kernel and bandwidth parameters as shown. The constructor accepts optional encoder module and dictionary parameters as well.

```python
from overcomplete import JumpSAE

# define a JumpReLU SAE with input dimension 768 and 10k concepts
sae = JumpSAE(768, 10_000)

# adjust kernel and bandwith
sae = JumpSAE(768, 10_000, bandwith = 1e-2,
              kernel='silverman')
```

--------------------------------

### RAJumpSAE.__init__

Source: https://kempnerinstitute.github.io/overcomplete/saes/archetypal_saes

Initializes a Relaxed Archetypal Jump SAE with the given input shape, number of concepts, archetype points, and bandwidth parameters.

```APIDOC
## RAJumpSAE.__init__

### Description
Initializes a Relaxed Archetypal Jump SAE.

### Method
Constructor

### Signature
`__init__(self, input_shape, nb_concepts, points, bandwidth=0.001, delta=1.0, use_multiplier=True, **kwargs)`

### Parameters
- **input_shape** (int) - Required - Dimensionality of the input data (excluding the batch dimension).
- **nb_concepts** (int) - Required - Number of dictionary atoms (concepts).
- **points** (torch.Tensor) - Required - The data points used to initialize/define the archetypes. Shape should be (num_points, input_shape).
- **bandwidth** (float) - Optional - Bandwidth parameter for the Jump SAE kernel, by default 1e-3.
- **delta** (float) - Optional - Delta parameter for the archetypal dictionary, by default 1.0.
- **use_multiplier** (bool) - Optional - Whether to use a learnable multiplier that parametrize the ball (e.g. if this parameter is 3 then the dictionary atoms are all on the ball of radius 3). By default True.
- **kwargs** (dict) - Optional - Additional arguments passed to the parent JumpSAE (e.g., encoder_module, device).
```

--------------------------------

### ConvexNMF.__init__

Source: https://kempnerinstitute.github.io/overcomplete/optimization/convexnmf

Initializes the ConvexNMF model with specified parameters.

```APIDOC
## ConvexNMF.__init__

### Description
Initializes the ConvexNMF model with specified parameters.

### Method
Constructor

### Parameters
- **nb_concepts** (int) - Required - Number of components to learn.
- **device** (str, optional) - Device to use for tensor computations, by default 'cpu'.
- **tol** (float, optional) - Tolerance value for the stopping criterion, by default 1e-4.
- **strict_convex** (bool, optional) - Whether to enforce the convexity constraint, by default False.
- **solver** (str, optional) - Optimization solver to use, either 'mu' (Multiplicative Update) or 'pgd' like method, by default 'mu'.
- **verbose** (bool, optional) - Whether to print optimization information, by default False.
- **l1_penalty** (float, optional) - L1 penalty coefficient, by default 0.0. Only used with the 'pgd' solver.

### Return
None
```

--------------------------------

### Compute Key Dictionary Learning Metrics

Source: https://kempnerinstitute.github.io/overcomplete/metrics

Demonstrates how to import and use common metrics from the overcomplete.metrics module to evaluate dictionary elements, structure, and reconstruction quality using PyTorch tensors.

```python
from overcomplete.metrics import (dead_codes, r2_score, hungarian_loss,
                                 cosine_hungarian_loss, wasserstein_1d)
# inputs
x = torch.randn(100, 10)
x_hat = torch.randn(100, 10)
# dictionaries
dict1 = torch.randn(512, 256)
dict2 = torch.randn(512, 256)
# concept values (codes)
codes = torch.randn(100, 512)
codes_2 = torch.randn(100, 512)

# check for inactive dictionary elements
dead_code_ratio = dead_codes(codes).mean()

# compare dictionary structures
hungarian_dist = hungarian_loss(dict1, dict2)
cosine_hungarian_dist = cosine_hungarian_loss(dict1, dict2)

# compute reconstruction quality
r2 = r2_score(x, x_hat)
# distrib. reconstruction quality
wasserstein_dist = wasserstein_1d(x, x_hat)
```

--------------------------------

### Basic Usage of ConvexNMF

Source: https://kempnerinstitute.github.io/overcomplete/optimization/convexnmf

Demonstrates initializing the ConvexNMF model, fitting it to data, and performing encoding and decoding operations.

```python
from overcomplete import ConvexNMF

# define a Convex NMF model with 10k concepts using
# the multiplicative update solver
convex_nmf = ConvexNMF(nb_concepts=10_000, solver='mu')

# fit the model to input data A
Z, D = convex_nmf.fit(A)

# encode new data
Z = convex_nmf.encode(A)
# decode (reconstruct) data from codes
A_hat = convex_nmf.decode(Z)
```

--------------------------------

### Instantiate RATopKSAE and RAJumpSAE

Source: https://kempnerinstitute.github.io/overcomplete/saes/archetypal_saes

Instantiate both archetypal SAEs with random points. Provide reasonably diverse points (e.g., k-means cluster centers) for stable archetypes. use_multiplier allows atoms to scale beyond the convex hull; set False to stay tighter. All standard training utilities (train_sae, custom losses) work unchanged.

```python
import torch
from overcomplete.sae import RATopKSAE, RAJumpSAE

points = torch.randn(2_000, 768)  # e.g. k-means centroids or sampled activations

ra_topk = RATopKSAE(
    input_shape=768,
    nb_concepts=10_000,
    points=points,
    top_k=20,
    delta=1.0,          # relaxation radius
    use_multiplier=True # learnable scaling of the archetypal hull
)

ra_jump = RAJumpSAE(
    input_shape=768,
    nb_concepts=10_000,
    points=points,
    bandwidth=1e-3,
    delta=1.5
)
```

--------------------------------

### MpSAE.__init__

Source: https://kempnerinstitute.github.io/overcomplete/saes/mp

Initializes a Matching Pursuit Sparse Autoencoder with specified input shape, number of concepts, and optional parameters.

```APIDOC
## MpSAE.__init__

### Description
Initializes a Matching Pursuit Sparse Autoencoder (MpSAE) with the given input shape and number of latent concepts.

### Method
Constructor

### Signature
`__init__(self, input_shape, nb_concepts, k=1, dropout=None, encoder_module='identity', dictionary_params=None, device='cpu')`

### Parameters
- **input_shape** (int or tuple of int) - Required - Dimensionality of the input data (excluding the batch dimension).
- **nb_concepts** (int) - Required - Number of latent dimensions (components) of the autoencoder.
- **k** (int, optional) - The number of matching pursuit iterations to perform (must be > 0).
- **dropout** (float, optional) - Probability of dropping a dictionary element at each iteration (range: 0.0 - 1.0). If None, no dropout is applied.
- **encoder_module** (nn.Module or str, optional) - Custom encoder module (or its registered name). If None, a default encoder is used.
- **dictionary_params** (dict, optional) - Parameters that will be passed to the dictionary layer. See DictionaryLayer for more details.
- **device** (str, optional) - Device on which to run the model (default is 'cpu').
```

--------------------------------

### Basic SAE instantiation with default and custom encoder

Source: https://kempnerinstitute.github.io/overcomplete/saes/vanilla

Shows two ways to create an SAE: a basic one with default linear encoding, and a more complex one using batch normalization in the encoder and L1 normalization for the dictionary.

```python
from overcomplete import SAE

# Define a basic SAE where input dimension is 768, with 10k concepts
# Using a simple linear encoding
sae = SAE(768, 10_000)

# Define a more complex SAE with batch normalization in the encoder
# The dictionary is normalized on the L1 ball instead of L2
sae = SAE(768, 10_000, encoder_module='mlp_bn_1',
          dictionary_params={'normalization': 'l1'})
```

--------------------------------

### ConvexNMF.init_semi_nmf

Source: https://kempnerinstitute.github.io/overcomplete/optimization/convexnmf

Initializes the Convex NMF model using Semi-NMF.

```APIDOC
## ConvexNMF.init_semi_nmf

### Description
Initializes the Convex NMF model using Semi-NMF.

### Method
Method

### Parameters
- **A** (torch.Tensor) - Required - Input tensor of shape (n_samples, n_features).
- **max_snmf_iters** (int, optional) - Maximum number of iterations for Semi-NMF, by default 100.

### Return
- **Z** (torch.Tensor) - Initialized codes tensor.
- **D** (torch.Tensor) - Initialized dictionary tensor.
```

--------------------------------

### OMPSAE Constructor

Source: https://kempnerinstitute.github.io/overcomplete/saes/omp

Initializes an OMPSAE instance with specified input shape, number of concepts, pursuit steps, NNLS iterations, dropout, encoder module, dictionary parameters, and device.

```APIDOC
## OMPSAE.__init__

### Description
Initializes the Orthogonal Matching Pursuit Sparse Autoencoder with the given configuration.

### Method
Constructor

### Parameters
- **input_shape** (int or tuple of int) - Required - Dimensionality of the input data (excluding batch dimension).
- **nb_concepts** (int) - Required - Number of latent components (atoms) in the dictionary.
- **k** (int, optional) - Default number of pursuit iterations (must be > 0).
- **dropout** (float, optional) - Dropout rate applied to dictionary atoms (range [0.0, 1.0]).
- **encoder_module** (str or nn.Module, optional) - Encoder module or name of registered encoder.
- **dictionary_params** (dict, optional) - Parameters passed to the dictionary layer.
- **device** (str, optional) - Device to run the model on (default is 'cpu').
- **max_iter** (int, optional) - Default number of NNLS iterations (default: 10).

### Example
```python
sae = OMPSAE(
    input_shape=512,
    nb_concepts=4096,
    k=4,
    max_iter=15,
    dropout=0.1,
    encoder_module="identity",
    device="cuda"
)
```
```

--------------------------------

### JumpSAE.__init__

Source: https://kempnerinstitute.github.io/overcomplete/saes/jumprelu

Initializes a JumpReLU Sparse Autoencoder with specified input shape, number of concepts, kernel, bandwidth, and optional custom encoder and dictionary parameters.

```APIDOC
## JumpSAE.__init__

### Description
Initializes a JumpReLU Sparse Autoencoder (SAE) with the given input shape, number of concepts, kernel, bandwidth, and optional custom encoder and dictionary parameters.

### Method
Constructor

### Signature
`__init__(self, input_shape, nb_concepts, kernel='silverman', bandwidth=0.001, encoder_module=None, dictionary_params=None, device='cpu')`

### Parameters
- **input_shape** (int or tuple of int) - Required - Dimensionality of the input data, do not include batch dimensions. Usually 1d (dim), 2d (seq length, dim) or 3d (dim, height, width).
- **nb_concepts** (int) - Required - Number of components/concepts in the dictionary. The dictionary is overcomplete if the number of concepts > in_dimensions.
- **kernel** (str, optional) - Kernel function to use in the JumpReLU activation, by default 'silverman'. Current options: 'rectangle', 'gaussian', 'triangular', 'cosine', 'epanechnikov', 'quartic', 'silverman', 'cauchy'.
- **bandwidth** (float, optional) - Bandwidth of the kernel, by default 1e-3.
- **encoder_module** (nn.Module or string, optional) - Custom encoder module, by default None. If None, a simple Linear + BatchNorm default encoder is used. If string, the name of the registered encoder module.
- **dictionary_params** (dict, optional) - Parameters that will be passed to the dictionary layer. See DictionaryLayer for more details.
- **device** (str, optional) - Device to run the model on, by default 'cpu'.

### Return
None
```

--------------------------------

### NMF.__init__

Source: https://kempnerinstitute.github.io/overcomplete/optimization/nmf

Constructor for the NMF model. Initializes the model with number of concepts and various hyperparameters such as solver, tolerance, and verbosity.

```APIDOC
## NMF.__init__

### Description
Initializes the NMF model with the specified number of concepts and hyperparameters.

### Method
Constructor

### Parameters
- **nb_concepts** (int) - Required - Number of components to learn.
- **device** (str, optional) - Default 'cpu' - Device to use for tensor computations, by default 'cpu'
- **solver** (str, optional) - Default 'hals' - Optimization algorithm to use. Can be one of 'hals', 'mu', 'pgd', 'anls'.
- **tol** (float, optional) - Default 0.0001 - Tolerance value for the stopping criterion.
- **verbose** (bool, optional) - Default False - Whether to display a progress bar.

### Request Example
{
  "nb_concepts": 10,
  "device": "cuda",
  "solver": "hals",
  "tol": 0.0001,
  "verbose": false
}

```

--------------------------------

### RelaxedArchetypalDictionary.__init__

Source: https://kempnerinstitute.github.io/overcomplete/saes/archetypal

Initializes the RelaxedArchetypalDictionary with input dimensions, number of concepts, candidate data points, and optional parameters.

```APIDOC
## RelaxedArchetypalDictionary.__init__

### Description
Initializes the RelaxedArchetypalDictionary with the given parameters.

### Method
Constructor

### Parameters
- **in_dimensions** (int) - Required - Dimensionality of the input data (e.g number of channels).
- **nb_concepts** (int) - Required - Number of components/concepts in the dictionary. The dictionary is overcomplete if the number of concepts > in_dimensions.
- **points** (tensors) - Required - Real data points (or point in the convex hull) used to find the candidates archetypes.
- **delta** (float, optional) - Optional - Constraint on the relaxation term, by default 1.0.
- **use_multiplier** (bool, optional) - Optional - Whether to train a positive scalar to multiply the dictionary after convex combination, making the dictionary in the conical hull (and not convex hull) of the data points, by default True.
- **device** (str, optional) - Optional - Device to run the model on ('cpu' or 'cuda'), by default 'cpu'.

### Request Example
```python
RelaxedArchetypalDictionary(
    in_dimensions=768,
    nb_concepts=10_000,
    points=points,
    delta=1.0,
    use_multiplier=True,
    device='cpu'
)
```
```

--------------------------------

### Initialize RelaxedArchetypalDictionary and assign to SAE

Source: https://kempnerinstitute.github.io/overcomplete/saes/archetypal

Initialize a BatchTopKSAE and replace its dictionary with a RelaxedArchetypalDictionary. The points tensor should contain candidate data points (e.g., sampled from the dataset; the paper recommends k-means). After assignment, the SAE can be trained normally.

```python
import torch
from overcomplete.sae.batchtopk_sae import BatchTopKSAE
from overcomplete.sae.archetypal_dictionary import RelaxedArchetypalDictionary

# initialize any sae
sae = BatchTopKSAE(768, 10_000, top_k=50)

# assume 'points' is a tensor of candidate data points (e.g., sampled from your dataset)
# the original paper recommend k-means
points = torch.randn(1000, 768)

# create our ra-sae
archetypal_dict = RelaxedArchetypalDictionary(
    in_dimensions=768,
    nb_concepts=10_000,
    points=points,
    delta=1.0,
)

# set the SAE's dictionary with the archetypal dictionary
sae.dictionary = archetypal_dict
# you can now train normally your sae
```

--------------------------------

### BatchTopKSAE.untied

Source: https://kempnerinstitute.github.io/overcomplete/saes/batchtopk

Creates a new encoder with weight from the current dictionary (or random init).

```APIDOC
## BatchTopKSAE.untied

### Description
Creates a new encoder with weight from the current dictionary (or random init).

### Method
untied(self, bias=False, copy_from_dictionary=True)

### Parameters
- **bias** (bool, optional) - Optional - Whether to include bias in encoder, by default False.
- **copy_from_dictionary** (bool, optional) - Optional - If True, initialize encoder with current dictionary weights, by default True.

### Return
- **self** - Returns self for method chaining.

### Example
```python
sae = sae.untied()
```
```

--------------------------------

### SemiNMF.init_random_d

Source: https://kempnerinstitute.github.io/overcomplete/optimization/seminmf

Initializes the dictionary D using matrix inversion.

```APIDOC
## SemiNMF.init_random_d

### Description
Initializes the dictionary D using matrix inversion.

### Method
init_random_d

### Parameters
- **A** (torch.Tensor) - Required - Input tensor of shape (batch_size, n_features).
- **Z** (torch.Tensor) - Required - Codes tensor of shape (batch_size, nb_concepts).

### Returns
- **D** (torch.Tensor) - Initialized dictionary tensor.
```

--------------------------------

### MpSAE.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/saes/mp

Returns the learned dictionary tensor.

```APIDOC
## MpSAE.get_dictionary

### Description
Returns the learned dictionary tensor.

### Method
get_dictionary

### Signature
`get_dictionary(self)`

### Return
- **torch.Tensor** - Learned dictionary tensor of shape (nb_components, input_size).
```

--------------------------------

### JumpSAE.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/saes/jumprelu

Returns the learned dictionary tensor.

```APIDOC
## JumpSAE.get_dictionary

### Description
Returns the learned dictionary tensor.

### Method
get_dictionary

### Signature
`get_dictionary(self)`

### Parameters
None

### Return
- **torch.Tensor** - Learned dictionary tensor of shape (nb_components, input_size).
```

--------------------------------

### BatchTopKSAE.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/saes/batchtopk

Returns the learned dictionary tensor.

```APIDOC
## BatchTopKSAE.get_dictionary

### Description
Returns the learned dictionary.

### Method
get_dictionary(self)

### Return
- **torch.Tensor** - Learned dictionary tensor of shape (nb_components, input_size).

### Example
```python
dictionary = sae.get_dictionary()
```
```

--------------------------------

### SemiNMF.init_random_z

Source: https://kempnerinstitute.github.io/overcomplete/optimization/seminmf

Initializes the codes Z using random values (can be negative).

```APIDOC
## SemiNMF.init_random_z

### Description
Initializes the codes Z using random values (can be negative).

### Method
init_random_z

### Parameters
- **A** (torch.Tensor) - Required - Input tensor of shape (batch_size, n_features).

### Returns
- **Z** (torch.Tensor) - Initialized codes tensor.
```

--------------------------------

### ConvexNMF.init_random_w

Source: https://kempnerinstitute.github.io/overcomplete/optimization/convexnmf

Initializes the coefficient matrix W with non-negative random values.

```APIDOC
## ConvexNMF.init_random_w

### Description
Initializes the coefficient matrix W with non-negative random values.

### Method
Method

### Parameters
- **A** (torch.Tensor) - Required - Input tensor of shape (n_samples, n_features).

### Return
- **W** (torch.Tensor) - Initialized coefficient tensor.
```

--------------------------------

### TopKSAE.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/saes/topk

Returns the learned dictionary tensor.

```APIDOC
## TopKSAE.get_dictionary

### Description
Returns the learned dictionary.

### Method
get_dictionary

### Signature
```python
get_dictionary()
```

### Returns
- **torch.Tensor** - Learned dictionary tensor of shape (nb_components, input_size).
```

--------------------------------

### SemiNMF.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/optimization/seminmf

Returns the learned dictionary components from Semi-NMF.

```APIDOC
## SemiNMF.get_dictionary

### Description
Returns the learned dictionary components from Semi-NMF.

### Method
get_dictionary

### Parameters
None

### Returns
- **torch.Tensor** - Dictionary components.
```

--------------------------------

### OMPSAE.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/saes/omp

Returns the learned dictionary tensor.

```APIDOC
## OMPSAE.get_dictionary

### Description
Returns the learned dictionary.

### Method
get_dictionary

### Return
- **torch.Tensor** - Learned dictionary tensor of shape (nb_components, input_size).
```

--------------------------------

### NMF.init_nndsvda

Source: https://kempnerinstitute.github.io/overcomplete/optimization/nmf

Initializes the dictionary and codes using the Sklearn NNDsvdA algorithm.

```APIDOC
## NMF.init_nndsvda

### Description
Initialize the dictionary and the code using Sklearn NNDsvdA algorithm.

### Parameters
- **A** (torch.Tensor): Required - Input tensor of shape (batch_size, n_features).

### Return
- **Z** (torch.Tensor): Initialized codes tensor.
- **D** (torch.Tensor): Initialized dictionary tensor.

```

--------------------------------

### NMF.init_random_z

Source: https://kempnerinstitute.github.io/overcomplete/optimization/nmf

Initializes the codes Z using non-negative random values.

```APIDOC
## NMF.init_random_z

### Description
Initialize the codes Z using non negative random values.

### Parameters
- **A** (torch.Tensor): Required - Input tensor of shape (batch_size, n_features).

### Return
- **Z** (torch.Tensor): Initialized codes tensor.

```

--------------------------------

### evidence_top_images

Source: https://kempnerinstitute.github.io/overcomplete/visualization

Function to visualize the top activating image for a concept and highlight the top activating pixels.

```APIDOC
## evidence_top_images(images, heatmaps, concept_id, percentiles=None)

### Description
Visualize the top activating image for a concept and highlight the top activating pixels.

### Parameters
- **images** (torch.Tensor or PIL.Image or np.ndarray) - Required - Batch of input images of shape (batch_size, channels, height, width).
- **heatmaps** (torch.Tensor or np.ndarray) - Required - Batch of heatmaps corresponding to the input images of shape (batch_size, height, width, num_concepts).
- **concept_id** (int) - Required - Index of the concept to visualize.
- **percentiles** (list of int, optional) - Optional - List of percentiles to highlight, by default None.
```

--------------------------------

### NMF.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/optimization/nmf

Returns the learned dictionary components from the NMF model.

```APIDOC
## NMF.get_dictionary

### Description
Return the learned dictionary components from NMF.

### Return
- **torch.Tensor**: Dictionary components.

```

--------------------------------

### NMF.init_random_d

Source: https://kempnerinstitute.github.io/overcomplete/optimization/nmf

Initializes the dictionary D using non-negative random values.

```APIDOC
## NMF.init_random_d

### Description
Initialize the dictionary D using non negative random values.

### Parameters
- **A** (torch.Tensor): Required - Input tensor of shape (batch_size, n_features).

### Return
- **D** (torch.Tensor): Initialized dictionary tensor.

```

--------------------------------

### SAE.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/saes/vanilla

Returns the learned dictionary tensor.

```APIDOC
## SAE.get_dictionary

### Description
Returns the learned dictionary tensor.

### Method
get_dictionary(self)

### Parameters
None

### Return
- **torch.Tensor** - Learned dictionary tensor of shape (nb_components, input_size).
```

--------------------------------

### Basic Usage of MpSAE

Source: https://kempnerinstitute.github.io/overcomplete/saes/mp

Instantiate a Matching Pursuit SAE with input dimension 512 and 4096 concepts. Set k to the number of pursuit steps and dropout to optionally mask atoms each step. The encode method returns the residual and sparse codes.

```python
from overcomplete import MpSAE

# define a Matching Pursuit SAE with input dimension 512, 4k concepts
sae = MpSAE(512, 4_096, k=4, dropout=0.1)

# k = number of pursuit steps, dropout optionally masks atoms each step
residual, codes = sae.encode(x)

```

--------------------------------

### TopKSAE.untied

Source: https://kempnerinstitute.github.io/overcomplete/saes/topk

Creates a new encoder with weights from the current dictionary or random initialization.

```APIDOC
## TopKSAE.untied

### Description
Creates a new encoder with weight from the current dictionary (or random init).

### Method
untied

### Signature
```python
untied(bias=False, copy_from_dictionary=True)
```

### Parameters
- **bias** (bool, optional) - Optional - Whether to include bias in encoder, by default False.
- **copy_from_dictionary** (bool, optional) - Optional - If True, initialize encoder with current dictionary weights, by default True.

### Returns
- **self** - Returns self for method chaining.
```

--------------------------------

### BatchTopKSAE.forward

Source: https://kempnerinstitute.github.io/overcomplete/saes/batchtopk

Performs a forward pass through the autoencoder, returning pre_codes, codes, and reconstructed input.

```APIDOC
## BatchTopKSAE.forward

### Description
Performs a forward pass through the autoencoder.

### Method
forward(self, x)

### Parameters
- **x** (torch.Tensor) - Required - Input tensor of shape (batch_size, input_size).

### Return
- **SAEOuput** - Return the pre_codes (z_pre), codes (z) and reconstructed input tensor (x_hat).

### Example
```python
z_pre, z, x_hat = sae(x)
```
```

--------------------------------

### overlay_top_heatmaps

Source: https://kempnerinstitute.github.io/overcomplete/visualization

Function to overlay a heatmap on the top activating image for a given concept.

```APIDOC
## overlay_top_heatmaps(images, heatmaps, concept_id, cmap=None, alpha=0.35)

### Description
Visualize the top activating image for a concept and overlay the associated heatmap.

### Parameters
- **images** (torch.Tensor or PIL.Image or np.ndarray) - Required - Batch of input images of shape (batch_size, channels, height, width).
- **heatmaps** (torch.Tensor or np.ndarray) - Required - Batch of heatmaps corresponding to the input images of shape (batch_size, height, width, num_concepts).
- **concept_id** (int) - Required - Index of the concept to visualize.
- **cmap** (str, optional) - Optional - Colormap for the heatmap, by default 'jet'.
- **alpha** (float, optional) - Optional - Transparency of the heatmap overlay, by default 0.35.
```

--------------------------------

### ConvexNMF.init_random_z

Source: https://kempnerinstitute.github.io/overcomplete/optimization/convexnmf

Initializes the codes Z using non negative random values.

```APIDOC
## ConvexNMF.init_random_z

### Description
Initializes the codes Z using non negative random values.

### Method
Method

### Parameters
- **A** (torch.Tensor) - Required - Input tensor of shape (batch_size, n_features).

### Return
- **Z** (torch.Tensor) - Initialized codes tensor.
```

--------------------------------

### MpSAE.untied

Source: https://kempnerinstitute.github.io/overcomplete/saes/mp

Creates a new encoder with weights from the current dictionary or random initialization.

```APIDOC
## MpSAE.untied

### Description
Creates a new encoder with weights from the current dictionary (or random init).

### Method
untied

### Signature
`untied(self, bias=False, copy_from_dictionary=True)`

### Parameters
- **bias** (bool, optional) - Whether to include bias in encoder, by default False.
- **copy_from_dictionary** (bool, optional) - If True, initialize encoder with current dictionary weights, by default True.

### Return
- **self** - Returns self for method chaining.
```

--------------------------------

### TopKSAE.forward

Source: https://kempnerinstitute.github.io/overcomplete/saes/topk

Performs a forward pass through the autoencoder, returning pre-codes, codes, and reconstructed input.

```APIDOC
## TopKSAE.forward

### Description
Performs a forward pass through the autoencoder.

### Method
forward

### Signature
```python
forward(x)
```

### Parameters
- **x** (torch.Tensor) - Required - Input tensor of shape (batch_size, input_size).

### Returns
- **SAEOuput** - Return the pre_codes (z_pre), codes (z) and reconstructed input tensor (x_hat).
```

--------------------------------

### TopKSAE with custom encoder module

Source: https://kempnerinstitute.github.io/overcomplete/saes/topk

Creates a TopKSAE with top_k=10 and specifies the encoder module as 'mlp_bn_1'. The comment notes that you can also pass your own encoder module directly.

```python
# adjust the encoder module (you can also)
# directly pass your own encoder module
sae = TopKSAE(768, 10_000, top_k=10,
              encoder_module='mlp_bn_1')
```

--------------------------------

### OMPSAE.untied

Source: https://kempnerinstitute.github.io/overcomplete/saes/omp

Creates a new encoder with weights from the current dictionary or random initialization.

```APIDOC
## OMPSAE.untied

### Description
Creates a new encoder with weight from the current dictionary (or random init).

### Method
untied

### Parameters
- **bias** (bool, optional) - Whether to include bias in encoder, by default False.
- **copy_from_dictionary** (bool, optional) - If True, initialize encoder with current dictionary weights, by default True.

### Return
- **self** - Returns self for method chaining.
```

--------------------------------

### NMF.encode

Source: https://kempnerinstitute.github.io/overcomplete/optimization/nmf

Encodes input activations into codes using the NMF model.

```APIDOC
## NMF.encode

### Description
Encode the input tensor (the activations) using NMF.

### Parameters
- **A** (torch.Tensor or Iterable): Required - Input tensor of shape (batch_size, n_features).
- **max_iter** (int, optional): Default 100 - Maximum number of iterations.
- **tol** (float, optional): Default the value set in the constructor - Tolerance value for the stopping criterion.

### Return
- **torch.Tensor**: Encoded features (the codes).

```

--------------------------------

### SAE.untied

Source: https://kempnerinstitute.github.io/overcomplete/saes/vanilla

Creates a new encoder with weights from the current dictionary (or random init).

```APIDOC
## SAE.untied

### Description
Creates a new encoder with weights from the current dictionary (or random init).

### Method
untied(self, bias=False, copy_from_dictionary=True)

### Parameters
- **bias** (bool, optional) - Optional - Whether to include bias in encoder, by default False.
- **copy_from_dictionary** (bool, optional) - Optional - If True, initialize encoder with current dictionary weights, by default True.

### Return
- **self** - Returns self for method chaining.
```

--------------------------------

### ConvexNMF.get_dictionary

Source: https://kempnerinstitute.github.io/overcomplete/optimization/convexnmf

Returns the learned dictionary components from Convex NMF.

```APIDOC
## ConvexNMF.get_dictionary

### Description
Returns the learned dictionary components from Convex NMF.

### Method
Method

### Parameters
None

### Return
- **torch.Tensor** - Dictionary components D = A W.
```

--------------------------------

### NMF.fit

Source: https://kempnerinstitute.github.io/overcomplete/optimization/nmf

Fits the NMF model to the input data.

```APIDOC
## NMF.fit

### Description
Fit the NMF model to the input data.

### Parameters
- **A** (torch.Tensor or Iterable): Required - Input tensor of shape (batch_size, n_features).
- **max_iter** (int, optional): Default 500 - Maximum number of iterations.

### Return
No explicit return value documented.

```

--------------------------------

### JumpSAE.untied

Source: https://kempnerinstitute.github.io/overcomplete/saes/jumprelu

Creates a new encoder with weights from the current dictionary (or random init).

```APIDOC
## JumpSAE.untied

### Description
Creates a new encoder with weights from the current dictionary (or random init).

### Method
untied

### Signature
`untied(self, bias=False, copy_from_dictionary=True)`

### Parameters
- **bias** (bool, optional) - Whether to include bias in encoder, by default False.
- **copy_from_dictionary** (bool, optional) - If True, initialize encoder with current dictionary weights, by default True.

### Return
- **self** - Returns self for method chaining.
```

--------------------------------

### MpSAE.train

Source: https://kempnerinstitute.github.io/overcomplete/saes/mp

Hook called when switching between training and evaluation mode, ensuring no dropout during evaluation.

```APIDOC
## MpSAE.train

### Description
Hook called when switching between training and evaluation mode. We use it to ensure no dropout is applied during evaluation.

### Method
train

### Signature
`train(self, mode=True)`

### Parameters
- **mode** (bool, optional) - Whether to set the model to training mode or not, by default True.
```
Tokens:
Raw