pyautoencoder 1.1.0__tar.gz → 1.1.2__tar.gz

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyautoencoder
-Version: 1.1.0
+Version: 1.1.2
 Summary: A Python package offering implementations of state-of-the-art autoencoder architectures in PyTorch.
 Author: Andrea Pollastro
 License: MIT
@@ -67,6 +67,7 @@ PyAutoencoder is designed to offer **simple and easy access to autoencoder frame
 **Currently implemented**:
 - Autoencoder (AE)
 - Variational Autoencoder (VAE)
+- Adaptive Group Variational Autoencoder (Ada-GVAE)
 
 ---
 
@@ -123,8 +124,8 @@ for x in dataloader:
     loss_results.objective.backward() # negative ELBO
     optimizer.step()
     # optional: log components
-    log_likelihood = loss_results.components["log_likelihood"]
-    kl_divergence = loss_results.components["kl_divergence"]
+    log_likelihood = loss_results.diagnostics["log_likelihood"]
+    kl_divergence = loss_results.diagnostics["kl_divergence"]
 ```
 
 ## Examples
@@ -156,3 +157,6 @@ If you use this package in academic work, please cite:
   publisher={Elsevier}
 }
 ```
+## Acknowledgments
+
+This work was funded by the PNRR MUR project PE0000013-FAIR (CUP: E63C25000630006).

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/README.md

@@ -39,6 +39,7 @@ PyAutoencoder is designed to offer **simple and easy access to autoencoder frame
 **Currently implemented**:
 - Autoencoder (AE)
 - Variational Autoencoder (VAE)
+- Adaptive Group Variational Autoencoder (Ada-GVAE)
 
 ---
 
@@ -95,8 +96,8 @@ for x in dataloader:
     loss_results.objective.backward() # negative ELBO
     optimizer.step()
     # optional: log components
-    log_likelihood = loss_results.components["log_likelihood"]
-    kl_divergence = loss_results.components["kl_divergence"]
+    log_likelihood = loss_results.diagnostics["log_likelihood"]
+    kl_divergence = loss_results.diagnostics["kl_divergence"]
 ```
 
 ## Examples
@@ -128,3 +129,6 @@ If you use this package in academic work, please cite:
   publisher={Elsevier}
 }
 ```
+## Acknowledgments
+
+This work was funded by the PNRR MUR project PE0000013-FAIR (CUP: E63C25000630006).

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/docs/source/api/models.rst

@@ -90,4 +90,14 @@ Variational Autoencoder
 
 .. autoclass:: pyautoencoder.variational.VAEOutput
    :members:
-   :no-index:
+   :no-index:
+
+Adaptive Group Variational Autoencoder
+---------------------------------------
+
+.. autoclass:: pyautoencoder.variational.AdaGVAE
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :exclude-members: build, _encode, _decode
+   :special-members: __init__

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/docs/source/architecture.rst

@@ -137,7 +137,7 @@ Supported likelihoods:
   
   .. math::
   
-      \text{NLL} = \frac{1}{2}[(x-\hat{x})^2 + \log(2\pi)]
+      \text{NLL} = \frac{1}{2}(x-\hat{x})^2
 
 - **Bernoulli** – Discrete/binary data (logits)
 

pyautoencoder-1.1.2/pyautoencoder/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '1.1.2'
+__version_tuple__ = version_tuple = (1, 1, 2)
+
+__commit_id__ = commit_id = 'g2d2766837'

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/pyautoencoder/loss/base.py

@@ -82,7 +82,7 @@ def log_likelihood(x: torch.Tensor,
       .. math::
 
           \log p(x \mid \hat{x}) =
-              -\tfrac{1}{2} \left[ (x - \hat{x})^2 + \log(2\pi) \right].
+              -\tfrac{1}{2} (x - \hat{x})^2.
 
       The output has the same shape as ``x``. Summing over feature dimensions
       gives per-sample log-likelihoods.
@@ -134,8 +134,7 @@ def log_likelihood(x: torch.Tensor,
     
     elif likelihood == LikelihoodType.GAUSSIAN:
         squared_error = (x_hat - x) ** 2
-        log_2pi = _get_log2pi(x)
-        return -0.5 * (squared_error + log_2pi)
+        return -0.5 * squared_error
     
     else:
         raise ValueError(f"Unsupported likelihood: {likelihood}")
@@ -144,8 +143,8 @@ def kl_divergence_diag_gaussian(
     mu_q: torch.Tensor, 
     log_var_q: torch.Tensor, 
     mu_p: Optional[torch.Tensor] = None, 
-    log_var_p: Optional[torch.Tensor] = None
-    ) -> torch.Tensor:
+    log_var_p: Optional[torch.Tensor] = None,
+    reduce_sum: bool = True) -> torch.Tensor:
     r"""Compute the KL divergence :math:`\mathrm{KL}(q(z \mid x) \,\|\, p(z))`
     between two diagonal Gaussian distributions.
 
@@ -175,6 +174,8 @@ def kl_divergence_diag_gaussian(
         Mean of the second distribution ``[B, D_z]``. Defaults to 0.
     log_var_p : torch.Tensor, optional
         Log-variance of the second distribution ``[B, D_z]``. Defaults to 0.
+    reduce_sum: bool, optional
+        Sum over the dimensions. Default to True
 
     Returns
     -------
@@ -194,5 +195,6 @@ def kl_divergence_diag_gaussian(
     
     term1 = log_var_p - log_var_q
     term2 = (var_q + (mu_q - mu_p).pow(2)) / var_p
-    
-    return 0.5 * torch.sum(term1 + term2 - 1, dim=-1)
+    if reduce_sum:
+        return 0.5 * torch.sum(term1 + term2 - 1, dim=-1)
+    return 0.5 * (term1 + term2 - 1)

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/pyautoencoder/variational/__init__.py

@@ -1,4 +1,5 @@
 from .vae import (
+    AdaGVAE,
     VAE,
     VAEDecodeOutput,
     VAEEncodeOutput,
@@ -6,6 +7,7 @@ from .vae import (
 )
 
 __all__ = [
+    'AdaGVAE',
     'VAE',
     'VAEDecodeOutput',
     'VAEEncodeOutput',

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/pyautoencoder/variational/stochastic_layers.py

@@ -118,20 +118,24 @@ class FullyFactorizedGaussian(nn.Module):
         if S < 1:
             raise ValueError("S must be >= 1.")
 
-        mu = self.mu(x)            # type: ignore       # [B, Dz]
-        log_var = self.log_var(x)  # type: ignore       # [B, Dz]
+        mu = self.mu(x)            # type: ignore               # [B, Dz]
+        log_var = self.log_var(x)  # type: ignore               # [B, Dz]
 
         if self.training:
-            std = torch.exp(0.5 * log_var)              # [B, Dz]
-            mu_e  = mu.unsqueeze(1).expand(-1, S, -1)   # [B, S, Dz]
-            std_e = std.unsqueeze(1).expand(-1, S, -1)  # [B, S, Dz]
-            eps = torch.randn_like(std_e)
-            z = mu_e + std_e * eps                      # [B, S, Dz]
+            z = self.reparametrize(mu=mu, log_var=log_var, S=S) # [B, S, Dz]
         else:
-            z = mu.unsqueeze(1).expand(-1, S, -1)       # [B, S, Dz]
+            z = mu.unsqueeze(1).expand(-1, S, -1)               # [B, S, Dz]
 
         return z, mu, log_var
     
+    def reparametrize(self, mu: torch.Tensor, log_var: torch.Tensor, S: int = 1):
+        std = torch.exp(0.5 * log_var)              # [B, Dz]
+        mu_e  = mu.unsqueeze(1).expand(-1, S, -1)   # [B, S, Dz]
+        std_e = std.unsqueeze(1).expand(-1, S, -1)  # [B, S, Dz]
+        eps = torch.randn_like(std_e)
+        z = mu_e + std_e * eps                      # [B, S, Dz]
+        return z
+    
     @property
     def built(self) -> bool:
         """Whether the module has been successfully built.

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/pyautoencoder/variational/vae.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 from dataclasses import dataclass
-from typing import Union, Dict
+from typing import Union, Tuple
 
 from ..loss.base import (
     LikelihoodType, 
@@ -318,3 +318,210 @@ class VAE(BaseAutoencoder):
                 'kl_divergence': kl_q_p.mean().item(),
             }
         )
+
+class AdaGVAE(VAE):
+    r"""Adaptive Group Variational Autoencoder (Ada-GVAE), from Locatello et al. (2020).
+
+    This class extends the VAE class and enables feature disentanglement in the latent space. 
+    For inference, use the .encode() and .decode() methods, as the forward method expects pairs of images, 
+    following the formulation introduced by Locatello et al.
+    """
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        decoder: nn.Module,
+        latent_dim: int,
+    ):
+        """Construct an AdaGVAE from an encoder, decoder, and latent size.
+
+        Notes
+        -----
+        The encoder and decoder are identical to those in a standard VAE. The adaptive
+        grouping mechanism is applied during the encoding step when processing paired inputs.
+
+        Parameters
+        ----------
+        encoder : nn.Module
+            Maps input ``x`` to a feature vector ``f(x)`` with shape ``[B, F]``.
+        decoder : nn.Module
+            Maps latent samples ``z`` to reconstructions ``x_hat``.
+        latent_dim : int
+            Dimensionality ``D_z`` of the latent space.
+        """
+        super().__init__(encoder=encoder, decoder=decoder, latent_dim=latent_dim)
+
+    # --- training-time hooks required by BaseAutoencoder ---
+    def _encode_pair(self, x1: torch.Tensor, x2: torch.Tensor, S: int = 1) -> Tuple[VAEEncodeOutput, VAEEncodeOutput]:
+        r"""Encode a pair of inputs with adaptive posterior alignment.
+
+        As described in Locatello et al., this method:
+
+        1. Encodes both inputs independently to obtain posterior parameters ``(mu1, log_var1)`` and ``(mu2, log_var2)``.
+        2. Computes element-wise KL divergence between the two posteriors: ``KL(q1||q2) → [B, D_z]``.
+        3. Computes a per-sample threshold ``tau`` based on KL divergences.
+        4. For each dimension, selects aligned (shared) or independent posteriors:
+           - If ``KL(q1_d||q2_d) < tau``: uses average distribution ``q_tilde``.
+           - If ``KL(q1_d||q2_d) ≥ tau``: uses original independent distribution.
+        5. Samples from the resulting (mixed) posteriors.
+
+        Parameters
+        ----------
+        x1 : torch.Tensor
+            First input batch of shape ``[B, ...]``.
+        x2 : torch.Tensor
+            Second input batch of shape ``[B, ...]``.
+        S : int
+            Number of latent samples per input.
+
+        Returns
+        -------
+        Tuple[VAEEncodeOutput, VAEEncodeOutput]
+            A pair of ``VAEEncodeOutput`` objects, each containing:
+
+            - ``z`` of shape ``[B, S, D_z]``: samples from the adapted posteriors.
+            - ``mu`` of shape ``[B, D_z]``: the (adapted) means.
+            - ``log_var`` of shape ``[B, D_z]``: the (adapted) log-variances.
+
+        Notes
+        -----
+        The thresholding mechanism promotes learning of shared latent factors
+        while allowing independent variation for high-divergence dimensions. 
+        This encourages disentanglement and structured representations.
+        """
+        _, mu1, log_var1 = self.sampling_layer(self.encoder(x1))
+        _, mu2, log_var2 = self.sampling_layer(self.encoder(x2))
+
+        # KL(q1||q2) -> [B, latents]
+        kl_q1_q2 = kl_divergence_diag_gaussian(mu1, log_var1, mu2, log_var2, reduce_sum=False)
+
+        # Computing threshold tau
+        max_delta = torch.max(kl_q1_q2, dim=1, keepdim=True)[0]
+        min_delta = torch.min(kl_q1_q2, dim=1, keepdim=True)[0]
+        tau = 0.5 * (max_delta + min_delta)
+
+        # Computing q_tilde1 and q_tilde2
+        mu_mean = 0.5*(mu1 + mu2)
+        var_mean = 0.5*(torch.exp(log_var1) + torch.exp(log_var2))
+        log_var_mean = torch.log(var_mean)
+
+        mask = kl_q1_q2 < tau
+        mu_tilde1 = torch.where(mask, mu_mean, mu1)
+        mu_tilde2 = torch.where(mask, mu_mean, mu2)
+        log_var_tilde1 = torch.where(mask, log_var_mean, log_var1)
+        log_var_tilde2 = torch.where(mask, log_var_mean, log_var2)
+
+        z1 = self.sampling_layer.reparametrize(mu=mu_tilde1, log_var=log_var_tilde1, S=S)
+        z2 = self.sampling_layer.reparametrize(mu=mu_tilde2, log_var=log_var_tilde2, S=S)
+
+        return VAEEncodeOutput(z=z1, mu=mu_tilde1, log_var=log_var_tilde1), \
+               VAEEncodeOutput(z=z2, mu=mu_tilde2, log_var=log_var_tilde2)
+
+
+    def forward(self, x1: torch.Tensor, x2: torch.Tensor, S: int = 1) -> Tuple[VAEOutput, VAEOutput]:
+        """Full AdaGVAE forward pass: encode pairs with adaptive grouping, sample, and decode.
+
+        Parameters
+        ----------
+        x1 : torch.Tensor
+            First input batch of shape ``[B, ...]``.
+        x2 : torch.Tensor
+            Second input batch of shape ``[B, ...]``.
+        S : int
+            Number of latent samples for Monte Carlo estimates.
+
+        Returns
+        -------
+        Tuple[VAEOutput, VAEOutput]
+            A pair of VAE outputs, each containing:
+
+            - ``x_hat``: reconstructions from the adapted latent samples.
+            - ``z``: latent samples from the adapted posteriors.
+            - ``mu``: (adapted) posterior means.
+            - ``log_var``: (adapted) posterior log-variances.
+        """
+        x1_enc, x2_enc = self._encode_pair(x1, x2, S=S)
+        x1_dec = self._decode(x1_enc.z)
+        x2_dec = self._decode(x2_enc.z)
+        return VAEOutput(x_hat=x1_dec.x_hat, z=x1_enc.z, mu=x1_enc.mu, log_var=x1_enc.log_var), \
+               VAEOutput(x_hat=x2_dec.x_hat, z=x2_enc.z, mu=x2_enc.mu, log_var=x2_enc.log_var)
+    
+    def compute_loss(self,
+                     x1: torch.Tensor, 
+                     x1_vae_output: VAEOutput,
+                     x2: torch.Tensor, 
+                     x2_vae_output: VAEOutput,
+                     beta: float = 1,
+                     likelihood: Union[str, LikelihoodType] = LikelihoodType.GAUSSIAN) -> LossResult:
+        r"""Compute the combined ELBO for a pair of inputs with adaptive posteriors.
+
+        This method computes the sum of the standard VAE ELBOs for both inputs:
+
+        .. math::
+
+            \mathcal{L}(x_1, x_2; \beta)
+                = \left[ \mathbb{E}_{q(\hat{z} \mid x_1)}[\log p(x_1 \mid \hat{z})]
+                \;-\; \beta \, \mathrm{KL}(q(\hat{z} \mid x_1) \,\|\, p(\hat{z})) \right]
+                + \left[ \mathbb{E}_{q(\hat{z} \mid x_2)}[\log p(x_2 \mid \hat{z})]
+                \;-\; \beta \, \mathrm{KL}(q(\hat{z} \mid x_2) \,\|\, p(\hat{z})) \right].
+
+        The key difference from standard VAE is that the posteriors :math:`q(\hat{z} | x_1)` and
+        :math:`q(\hat{z} | x_2)` are obtained from the adaptive grouping mechanism, which can
+        share dimensions based on KL divergence thresholds.
+
+        Parameters
+        ----------
+        x1 : torch.Tensor
+            First input batch of shape ``[B, ...]``.
+        x1_vae_output : VAEOutput
+            Output from the forward pass for ``x1``. Expected fields:
+
+            - ``x_hat`` (torch.Tensor): Reconstructions, shape ``[B, ...]`` or ``[B, S, ...]``.
+            - ``mu`` (torch.Tensor): (Adapted) posterior mean, shape ``[B, D_z]``.
+            - ``log_var`` (torch.Tensor): (Adapted) posterior log-variance, shape ``[B, D_z]``.
+
+        x2 : torch.Tensor
+            Second input batch of shape ``[B, ...]``.
+        x2_vae_output : VAEOutput
+            Output from the forward pass for ``x2``. Same structure as ``x1_vae_output``.
+        likelihood : Union[str, LikelihoodType], optional
+            Likelihood model for the reconstruction term.
+            Can be 'gaussian' or 'bernoulli'. Defaults to Gaussian.
+        beta : float, optional
+            Weighting factor for the KL term (beta-VAE).
+            ``beta = 1`` yields the standard objective. Defaults to 1.
+
+        Returns
+        -------
+        LossResult
+            Result containing:
+
+            * **objective** – Sum of negative ELBOs for both inputs (scalar).
+            * **diagnostics** – Dictionary with:
+
+              - ``"elbo"``: Sum of ELBOs for both inputs.
+              - ``"log_likelihood_x1"``: Mean reconstruction term for ``x1``.
+              - ``"log_likelihood_x2"``: Mean reconstruction term for ``x2``.
+              - ``"kl_divergence_x1"``: Mean KL divergence for ``x1``'s posterior.
+              - ``"kl_divergence_x2"``: Mean KL divergence for ``x2``'s posterior.
+
+        Notes
+        -----
+        - All diagnostics are **batch means** (per-sample losses averaged over ``B``).
+        - Gradients flow through both decoders; neither input is detached.
+        - The adaptive grouping introduces implicit structure learning through
+          the selective sharing of posterior dimensions.
+        """
+        x1_loss_info = super().compute_loss(x=x1, vae_output=x1_vae_output, beta=beta, likelihood=likelihood)
+        x2_loss_info = super().compute_loss(x=x2, vae_output=x2_vae_output, beta=beta, likelihood=likelihood)
+
+        return LossResult(
+            objective = x1_loss_info.objective + x2_loss_info.objective,
+            diagnostics = {
+                'elbo': x1_loss_info.diagnostics['elbo'] + x2_loss_info.diagnostics['elbo'],
+                'log_likelihood_x1': x1_loss_info.diagnostics['log_likelihood'],
+                'log_likelihood_x2': x2_loss_info.diagnostics['log_likelihood'],
+                'kl_divergence_x1': x1_loss_info.diagnostics['kl_divergence'],
+                'kl_divergence_x2': x2_loss_info.diagnostics['kl_divergence'],
+            }
+        )

{pyautoencoder-1.1.0 → pyautoencoder-1.1.2}/pyautoencoder.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyautoencoder
-Version: 1.1.0
+Version: 1.1.2
 Summary: A Python package offering implementations of state-of-the-art autoencoder architectures in PyTorch.
 Author: Andrea Pollastro
 License: MIT
@@ -67,6 +67,7 @@ PyAutoencoder is designed to offer **simple and easy access to autoencoder frame
 **Currently implemented**:
 - Autoencoder (AE)
 - Variational Autoencoder (VAE)
+- Adaptive Group Variational Autoencoder (Ada-GVAE)
 
 ---
 
@@ -123,8 +124,8 @@ for x in dataloader:
     loss_results.objective.backward() # negative ELBO
     optimizer.step()
     # optional: log components
-    log_likelihood = loss_results.components["log_likelihood"]
-    kl_divergence = loss_results.components["kl_divergence"]
+    log_likelihood = loss_results.diagnostics["log_likelihood"]
+    kl_divergence = loss_results.diagnostics["kl_divergence"]
 ```
 
 ## Examples
@@ -156,3 +157,6 @@ If you use this package in academic work, please cite:
   publisher={Elsevier}
 }
 ```
+## Acknowledgments
+
+This work was funded by the PNRR MUR project PE0000013-FAIR (CUP: E63C25000630006).