Skip to content

naive_probabilistic_compression

Bases: transformation

A probabilistic-based compression class for dimensionality reduction.

This class compresses input data by sampling features probabilistically based on a specified distribution or metric. It supports simple sampling, normalization, and probabilistic weighting.

Notes

Formally, given a data instance \(\mathbf{x} \in {R}^m\), we define the probabilistic compression function based on probabilistic sampling as:

$$ \begin{equation} \kappa(\mathbf{x}) = \mathbf{t} \in {R}^d, \end{equation} $$ where the output vector \(\mathbf{t}\) is conditionally dependent on \(\mathbf{x}\) following certain distributions. For example, using a Gaussian distribution:

$$ \begin{equation} \mathbf{t} | \mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma}). \end{equation} $$ The dimension \(d\) of the output vector \(\mathbf{t}\) is a hyper-parameter \(d = k\) requiring manual setup.

Attributes:

Name Type Description
k int

Number of features to retain after compression.
metric (Callable, optional)

Metric function to apply to the input tensor before sampling. Defaults to None.
simply_sampling bool

If True, performs simple sampling without further processing. Defaults to True.
with_replacement bool

If True, samples features with replacement. Defaults to False.
require_normalization bool

If True, normalizes the input tensor before sampling. Defaults to True.
log_prob bool

If True, returns the logarithm of probabilities for the compressed features. Defaults to False.
distribution_function distributions

Probability distribution function used for sampling. Defaults to a uniform distribution.

Methods:

Name Description
__init__

Initializes the probabilistic compression function.
calculate_D

Validates and returns the number of features to retain (k).
to_config

Converts the current configuration into a dictionary format.
calculate_weights

Computes sampling weights for the input tensor based on the probability distribution.
forward

Applies probabilistic sampling to compress the input tensor.
Source code in tinybig/compression/probabilistic_compression.py 
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
class naive_probabilistic_compression(transformation):
    r"""
        A probabilistic-based compression class for dimensionality reduction.

        This class compresses input data by sampling features probabilistically based on a specified
        distribution or metric. It supports simple sampling, normalization, and probabilistic weighting.

        Notes
        ----------
        Formally, given a data instance $\mathbf{x} \in {R}^m$, we define the probabilistic compression function based on probabilistic sampling as:

        $$
            \begin{equation}
            \kappa(\mathbf{x}) = \mathbf{t} \in {R}^d,
            \end{equation}
        $$
        where the output vector $\mathbf{t}$ is conditionally dependent on $\mathbf{x}$ following certain distributions. For example, using a Gaussian distribution:

        $$
            \begin{equation}
            \mathbf{t} | \mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma}).
            \end{equation}
        $$
        The dimension $d$ of the output vector $\mathbf{t}$ is a hyper-parameter $d = k$ requiring manual setup.

        Attributes
        ----------
        k : int
            Number of features to retain after compression.
        metric : Callable, optional
            Metric function to apply to the input tensor before sampling. Defaults to None.
        simply_sampling : bool
            If True, performs simple sampling without further processing. Defaults to True.
        with_replacement : bool
            If True, samples features with replacement. Defaults to False.
        require_normalization : bool
            If True, normalizes the input tensor before sampling. Defaults to True.
        log_prob : bool
            If True, returns the logarithm of probabilities for the compressed features. Defaults to False.
        distribution_function : torch.distributions
            Probability distribution function used for sampling. Defaults to a uniform distribution.

        Methods
        -------
        __init__(k, name='probabilistic_compression', simply_sampling=True, distribution_function=None, ...)
            Initializes the probabilistic compression function.
        calculate_D(m: int)
            Validates and returns the number of features to retain (`k`).
        to_config()
            Converts the current configuration into a dictionary format.
        calculate_weights(x: torch.Tensor)
            Computes sampling weights for the input tensor based on the probability distribution.
        forward(x: torch.Tensor, device='cpu', *args, **kwargs)
            Applies probabilistic sampling to compress the input tensor.
    """
    def __init__(
        self,
        k: int,
        name: str = 'probabilistic_compression',
        simply_sampling: bool = True,
        distribution_function: torch.distributions = None,
        distribution_function_configs: dict = None,
        metric: Callable[[torch.Tensor], torch.Tensor] = None,
        with_replacement: bool = False,
        require_normalization: bool = True,
        log_prob: bool = False,
        *args, **kwargs
    ):
        """
            The initialization method of the naive probabilistic compression function.

            It initializes the compression function based on the provided probabilistic distribution function.

            Parameters
            ----------
            k : int
                Number of features to retain after compression.
            name : str, optional
                Name of the transformation. Defaults to 'probabilistic_compression'.
            simply_sampling : bool, optional
                If True, performs simple sampling without further processing. Defaults to True.
            distribution_function : torch.distributions, optional
                Pre-defined probability distribution function for sampling. Defaults to None.
            distribution_function_configs : dict, optional
                Configuration dictionary for initializing the distribution function. Defaults to None.
            metric : Callable, optional
                Metric function to apply to the input tensor before sampling. Defaults to None.
            with_replacement : bool, optional
                If True, samples features with replacement. Defaults to False.
            require_normalization : bool, optional
                If True, normalizes the input tensor before sampling. Defaults to True.
            log_prob : bool, optional
                If True, returns the logarithm of probabilities for the compressed features. Defaults to False.
            *args : tuple
                Additional positional arguments for the parent `transformation` class.
            **kwargs : dict
                Additional keyword arguments for the parent `transformation` class.
        """

        super().__init__(name=name, *args, **kwargs)
        self.k = k
        self.metric = metric
        self.simply_sampling = simply_sampling
        self.with_replacement = with_replacement
        self.require_normalization = require_normalization
        self.log_prob = log_prob

        if self.simply_sampling:
            self. log_prob = False

        if distribution_function is not None:
            self.distribution_function = distribution_function
        elif distribution_function_configs is not None:
            function_class = distribution_function_configs['function_class']
            function_parameters = distribution_function_configs['function_parameters'] if 'function_parameters' in distribution_function_configs else {}
            self.distribution_function = config.get_obj_from_str(function_class)(**function_parameters)
        else:
            self.distribution_function = None

        if self.distribution_function is None:
            self.distribution_function = torch.distributions.uniform.Uniform(low=0.0, high=1.0)

    def calculate_D(self, m: int):
        """
            Validates and returns the number of features to retain (`k`).

            This method ensures that the number of features to retain (`k`) is within the valid range
            [0, m], where `m` is the total number of features in the input.

            Parameters
            ----------
            m : int
                Total number of features in the input tensor.

            Returns
            -------
            int
                The number of features to retain (`k`).

            Raises
            ------
            AssertionError
                If `k` is not set or is not within the range [0, m].
        """
        assert self.k is not None and 0 <= self.k <= m
        return self.k

    def to_config(self):
        """
            Converts the current configuration into a dictionary format.

            This method extracts the current configuration of the instance, including the distribution
            function and its parameters, and returns it as a dictionary.

            Returns
            -------
            dict
                A dictionary containing the configuration of the instance, including parameters for the
                distribution function.
        """
        configs = super().to_config()
        configs['function_parameters'].pop('distribution_function')
        if self.distribution_function is not None:
            configs['function_parameters']['distribution_function_configs'] = function.functions_to_configs(self.distribution_function)
        return configs

    def calculate_weights(self, x: torch.Tensor):
        """
            Computes sampling weights for the input tensor based on the probability distribution.

            This method applies the specified distribution function to compute weights for sampling.
            If no distribution function is provided, uniform weights are assigned.

            Parameters
            ----------
            x : torch.Tensor
                Input tensor of shape `(batch_size, num_features)`.

            Returns
            -------
            torch.Tensor
                Sampling weights of shape `(batch_size, num_features)`, normalized to sum to 1 along
                the last dimension.
        """
        if self.distribution_function is not None:
            x = torch.exp(self.distribution_function.log_prob(x))
            weights = x/x.sum(dim=-1, keepdim=True)
        else:
            b, m = x.shape
            weights = torch.ones((b, m)) / m
        return weights

    def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
        r"""
            Applies probabilistic sampling to compress the input tensor.

            This method processes the input tensor, computes sampling weights, and uses them to select
            features probabilistically based on the specified distribution function or metric.

            Formally, given a data instance $\mathbf{x} \in {R}^m$, we define the probabilistic compression function based on probabilistic sampling as:

            $$
                \begin{equation}
                \kappa(\mathbf{x}) = \mathbf{t} \in {R}^d,
                \end{equation}
            $$
            where the output vector $\mathbf{t}$ is conditionally dependent on $\mathbf{x}$ following certain distributions. For example, using a Gaussian distribution:

            $$
                \begin{equation}
                \mathbf{t} | \mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma}).
                \end{equation}
            $$
            The dimension $d$ of the output vector $\mathbf{t}$ is a hyper-parameter $d = k$ requiring manual setup.


            Parameters
            ----------
            x : torch.Tensor
                Input tensor of shape `(batch_size, num_features)`.
            device : str, optional
                Device for computation (e.g., 'cpu' or 'cuda'). Defaults to 'cpu'.
            *args : tuple
                Additional positional arguments for pre- and post-processing.
            **kwargs : dict
                Additional keyword arguments for pre- and post-processing.

            Returns
            -------
            torch.Tensor
                Compressed tensor of shape `(batch_size, k)`.

            Raises
            ------
            AssertionError
                If the output tensor shape does not match the expected `(batch_size, k)`.
        """
        b, m = x.shape
        x = self.pre_process(x=x, device=device)

        data_x = None
        if self.simply_sampling:
            data_x = x.clone()

        if self.metric is not None:
            x = self.metric(x)
        if self.require_normalization:
            x = 0.99 * torch.nn.functional.sigmoid(x) + 0.001

        weights = self.calculate_weights(x)
        sampled_indices = torch.multinomial(weights, self.calculate_D(m=m), replacement=self.with_replacement)
        sampled_indices, _ = torch.sort(sampled_indices, dim=1)

        if self.simply_sampling:
            compression = torch.gather(data_x, 1, sampled_indices)
        else:
            compression = torch.gather(x, 1, sampled_indices)

        if self.log_prob:
            compression = self.distribution_function.log_prob(compression)

        assert compression.shape == (b, self.calculate_D(m=m))
        return self.post_process(x=compression, device=device)

__init__(k, name='probabilistic_compression', simply_sampling=True, distribution_function=None, distribution_function_configs=None, metric=None, with_replacement=False, require_normalization=True, log_prob=False, *args, **kwargs)

The initialization method of the naive probabilistic compression function.

It initializes the compression function based on the provided probabilistic distribution function.

Parameters:

Name Type Description Default
k int

Number of features to retain after compression.

 required
name str

Name of the transformation. Defaults to 'probabilistic_compression'.

 'probabilistic_compression'
simply_sampling bool

If True, performs simple sampling without further processing. Defaults to True.

 True
distribution_function distributions

Pre-defined probability distribution function for sampling. Defaults to None.

 None
distribution_function_configs dict

Configuration dictionary for initializing the distribution function. Defaults to None.

 None
metric Callable

Metric function to apply to the input tensor before sampling. Defaults to None.

 None
with_replacement bool

If True, samples features with replacement. Defaults to False.

 False
require_normalization bool

If True, normalizes the input tensor before sampling. Defaults to True.

 True
log_prob bool

If True, returns the logarithm of probabilities for the compressed features. Defaults to False.

 False
*args tuple

Additional positional arguments for the parent transformation class.

 ()
**kwargs dict

Additional keyword arguments for the parent transformation class.

 {}
Source code in tinybig/compression/probabilistic_compression.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
def __init__(
    self,
    k: int,
    name: str = 'probabilistic_compression',
    simply_sampling: bool = True,
    distribution_function: torch.distributions = None,
    distribution_function_configs: dict = None,
    metric: Callable[[torch.Tensor], torch.Tensor] = None,
    with_replacement: bool = False,
    require_normalization: bool = True,
    log_prob: bool = False,
    *args, **kwargs
):
    """
        The initialization method of the naive probabilistic compression function.

        It initializes the compression function based on the provided probabilistic distribution function.

        Parameters
        ----------
        k : int
            Number of features to retain after compression.
        name : str, optional
            Name of the transformation. Defaults to 'probabilistic_compression'.
        simply_sampling : bool, optional
            If True, performs simple sampling without further processing. Defaults to True.
        distribution_function : torch.distributions, optional
            Pre-defined probability distribution function for sampling. Defaults to None.
        distribution_function_configs : dict, optional
            Configuration dictionary for initializing the distribution function. Defaults to None.
        metric : Callable, optional
            Metric function to apply to the input tensor before sampling. Defaults to None.
        with_replacement : bool, optional
            If True, samples features with replacement. Defaults to False.
        require_normalization : bool, optional
            If True, normalizes the input tensor before sampling. Defaults to True.
        log_prob : bool, optional
            If True, returns the logarithm of probabilities for the compressed features. Defaults to False.
        *args : tuple
            Additional positional arguments for the parent `transformation` class.
        **kwargs : dict
            Additional keyword arguments for the parent `transformation` class.
    """

    super().__init__(name=name, *args, **kwargs)
    self.k = k
    self.metric = metric
    self.simply_sampling = simply_sampling
    self.with_replacement = with_replacement
    self.require_normalization = require_normalization
    self.log_prob = log_prob

    if self.simply_sampling:
        self. log_prob = False

    if distribution_function is not None:
        self.distribution_function = distribution_function
    elif distribution_function_configs is not None:
        function_class = distribution_function_configs['function_class']
        function_parameters = distribution_function_configs['function_parameters'] if 'function_parameters' in distribution_function_configs else {}
        self.distribution_function = config.get_obj_from_str(function_class)(**function_parameters)
    else:
        self.distribution_function = None

    if self.distribution_function is None:
        self.distribution_function = torch.distributions.uniform.Uniform(low=0.0, high=1.0)

calculate_D(m)

Validates and returns the number of features to retain (k).

This method ensures that the number of features to retain (k) is within the valid range [0, m], where m is the total number of features in the input.

Parameters:

Name Type Description Default
m int

Total number of features in the input tensor.

 required

Returns:

Type Description
int

The number of features to retain (k).

Raises:

Type Description
AssertionError

If k is not set or is not within the range [0, m].
Source code in tinybig/compression/probabilistic_compression.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def calculate_D(self, m: int):
    """
        Validates and returns the number of features to retain (`k`).

        This method ensures that the number of features to retain (`k`) is within the valid range
        [0, m], where `m` is the total number of features in the input.

        Parameters
        ----------
        m : int
            Total number of features in the input tensor.

        Returns
        -------
        int
            The number of features to retain (`k`).

        Raises
        ------
        AssertionError
            If `k` is not set or is not within the range [0, m].
    """
    assert self.k is not None and 0 <= self.k <= m
    return self.k

calculate_weights(x)

Computes sampling weights for the input tensor based on the probability distribution.

This method applies the specified distribution function to compute weights for sampling. If no distribution function is provided, uniform weights are assigned.

Parameters:

Name Type Description Default
x Tensor

Input tensor of shape (batch_size, num_features).

 required

Returns:

Type Description
Tensor

Sampling weights of shape (batch_size, num_features), normalized to sum to 1 along the last dimension.
Source code in tinybig/compression/probabilistic_compression.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
def calculate_weights(self, x: torch.Tensor):
    """
        Computes sampling weights for the input tensor based on the probability distribution.

        This method applies the specified distribution function to compute weights for sampling.
        If no distribution function is provided, uniform weights are assigned.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape `(batch_size, num_features)`.

        Returns
        -------
        torch.Tensor
            Sampling weights of shape `(batch_size, num_features)`, normalized to sum to 1 along
            the last dimension.
    """
    if self.distribution_function is not None:
        x = torch.exp(self.distribution_function.log_prob(x))
        weights = x/x.sum(dim=-1, keepdim=True)
    else:
        b, m = x.shape
        weights = torch.ones((b, m)) / m
    return weights

forward(x, device='cpu', *args, **kwargs)

Applies probabilistic sampling to compress the input tensor.

This method processes the input tensor, computes sampling weights, and uses them to select features probabilistically based on the specified distribution function or metric.

Formally, given a data instance \(\mathbf{x} \in {R}^m\), we define the probabilistic compression function based on probabilistic sampling as:

$$ \begin{equation} \kappa(\mathbf{x}) = \mathbf{t} \in {R}^d, \end{equation} $$ where the output vector \(\mathbf{t}\) is conditionally dependent on \(\mathbf{x}\) following certain distributions. For example, using a Gaussian distribution:

$$ \begin{equation} \mathbf{t} | \mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma}). \end{equation} $$ The dimension \(d\) of the output vector \(\mathbf{t}\) is a hyper-parameter \(d = k\) requiring manual setup.

Parameters:

Name Type Description Default
x Tensor

Input tensor of shape (batch_size, num_features).

 required
device str

Device for computation (e.g., 'cpu' or 'cuda'). Defaults to 'cpu'.

 'cpu'
*args tuple

Additional positional arguments for pre- and post-processing.

 ()
**kwargs dict

Additional keyword arguments for pre- and post-processing.

 {}

Returns:

Type Description
Tensor

Compressed tensor of shape (batch_size, k).

Raises:

Type Description
AssertionError

If the output tensor shape does not match the expected (batch_size, k).
Source code in tinybig/compression/probabilistic_compression.py 
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
def forward(self, x: torch.Tensor, device='cpu', *args, **kwargs):
    r"""
        Applies probabilistic sampling to compress the input tensor.

        This method processes the input tensor, computes sampling weights, and uses them to select
        features probabilistically based on the specified distribution function or metric.

        Formally, given a data instance $\mathbf{x} \in {R}^m$, we define the probabilistic compression function based on probabilistic sampling as:

        $$
            \begin{equation}
            \kappa(\mathbf{x}) = \mathbf{t} \in {R}^d,
            \end{equation}
        $$
        where the output vector $\mathbf{t}$ is conditionally dependent on $\mathbf{x}$ following certain distributions. For example, using a Gaussian distribution:

        $$
            \begin{equation}
            \mathbf{t} | \mathbf{x} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\Sigma}).
            \end{equation}
        $$
        The dimension $d$ of the output vector $\mathbf{t}$ is a hyper-parameter $d = k$ requiring manual setup.


        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape `(batch_size, num_features)`.
        device : str, optional
            Device for computation (e.g., 'cpu' or 'cuda'). Defaults to 'cpu'.
        *args : tuple
            Additional positional arguments for pre- and post-processing.
        **kwargs : dict
            Additional keyword arguments for pre- and post-processing.

        Returns
        -------
        torch.Tensor
            Compressed tensor of shape `(batch_size, k)`.

        Raises
        ------
        AssertionError
            If the output tensor shape does not match the expected `(batch_size, k)`.
    """
    b, m = x.shape
    x = self.pre_process(x=x, device=device)

    data_x = None
    if self.simply_sampling:
        data_x = x.clone()

    if self.metric is not None:
        x = self.metric(x)
    if self.require_normalization:
        x = 0.99 * torch.nn.functional.sigmoid(x) + 0.001

    weights = self.calculate_weights(x)
    sampled_indices = torch.multinomial(weights, self.calculate_D(m=m), replacement=self.with_replacement)
    sampled_indices, _ = torch.sort(sampled_indices, dim=1)

    if self.simply_sampling:
        compression = torch.gather(data_x, 1, sampled_indices)
    else:
        compression = torch.gather(x, 1, sampled_indices)

    if self.log_prob:
        compression = self.distribution_function.log_prob(compression)

    assert compression.shape == (b, self.calculate_D(m=m))
    return self.post_process(x=compression, device=device)

to_config()

Converts the current configuration into a dictionary format.

This method extracts the current configuration of the instance, including the distribution function and its parameters, and returns it as a dictionary.

Returns:

Type Description
dict

A dictionary containing the configuration of the instance, including parameters for the distribution function.
Source code in tinybig/compression/probabilistic_compression.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def to_config(self):
    """
        Converts the current configuration into a dictionary format.

        This method extracts the current configuration of the instance, including the distribution
        function and its parameters, and returns it as a dictionary.

        Returns
        -------
        dict
            A dictionary containing the configuration of the instance, including parameters for the
            distribution function.
    """
    configs = super().to_config()
    configs['function_parameters'].pop('distribution_function')
    if self.distribution_function is not None:
        configs['function_parameters']['distribution_function_configs'] = function.functions_to_configs(self.distribution_function)
    return configs