Skip to content

potential_based

PotentialBased(input_size, hidden_size, activation_nonlin, tau_init, tau_learnable, kappa_init, kappa_learnable, potentials_init=None, output_size=None, input_embedding=None, output_embedding=None, device='cpu')

Bases: Module, ABC

Base class for all potential-based recurrent neutral networks.

hidden_size: Number of neurons with potential per hidden layer. For all use cases conceived at this point,
    we only use one recurrent layer. However, there is the possibility to extend the networks to multiple
    potential-based layers.
activation_nonlin: Nonlinearity used to compute the activations from the potential levels.
tau_init: Initial value for the shared time constant of the potentials.
tau_learnable: Whether the time constant is a learnable parameter or fixed.
kappa_init: Initial value for the cubic decay, pass 0 to disable the cubic decay.
kappa_learnable: Whether the cubic decay is a learnable parameter or fixed.
potentials_init: Initial for the potentials, i.e., the network's hidden state.
output_size: Number of output dimensions. By default, the number of outputs is equal to the number of
    hidden neurons.
input_embedding: Optional (custom) [Module][torch.nn.Module] to extract features from the inputs.
    This module must transform the inputs such that the dimensionality matches the number of
    neurons of the neural field, i.e., `hidden_size`. By default, a [linear layer][torch.nn.Linear]
    without biases is used.
output_embedding: Optional (custom) [Module][torch.nn.Module] to compute the outputs from the activations.
    This module must map the activations of shape (`hidden_size`,) to the outputs of shape (`output_size`,)
    By default, a [linear layer][torch.nn.Linear] without biases is used.
device: Device to move this module to (after initialization).
Source code in neuralfields/potential_based.py
 28
 29
 30
 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
def __init__(
    self,
    input_size: int,
    hidden_size: int,
    activation_nonlin: Union[ActivationFunction, Sequence[ActivationFunction]],
    tau_init: Union[float, int],
    tau_learnable: bool,
    kappa_init: Union[float, int],
    kappa_learnable: bool,
    potentials_init: Optional[torch.Tensor] = None,
    output_size: Optional[int] = None,
    input_embedding: Optional[nn.Module] = None,
    output_embedding: Optional[nn.Module] = None,
    device: Union[str, torch.device] = "cpu",
):
    """
    Args:
        input_size: Number of input dimensions.
        hidden_size: Number of neurons with potential per hidden layer. For all use cases conceived at this point,
            we only use one recurrent layer. However, there is the possibility to extend the networks to multiple
            potential-based layers.
        activation_nonlin: Nonlinearity used to compute the activations from the potential levels.
        tau_init: Initial value for the shared time constant of the potentials.
        tau_learnable: Whether the time constant is a learnable parameter or fixed.
        kappa_init: Initial value for the cubic decay, pass 0 to disable the cubic decay.
        kappa_learnable: Whether the cubic decay is a learnable parameter or fixed.
        potentials_init: Initial for the potentials, i.e., the network's hidden state.
        output_size: Number of output dimensions. By default, the number of outputs is equal to the number of
            hidden neurons.
        input_embedding: Optional (custom) [Module][torch.nn.Module] to extract features from the inputs.
            This module must transform the inputs such that the dimensionality matches the number of
            neurons of the neural field, i.e., `hidden_size`. By default, a [linear layer][torch.nn.Linear]
            without biases is used.
        output_embedding: Optional (custom) [Module][torch.nn.Module] to compute the outputs from the activations.
            This module must map the activations of shape (`hidden_size`,) to the outputs of shape (`output_size`,)
            By default, a [linear layer][torch.nn.Linear] without biases is used.
        device: Device to move this module to (after initialization).
    """
    # Call torch.nn.Module's constructor.
    super().__init__()

    # For all use cases conceived at this point, we only use one recurrent layer. However, this variable still
    # exists in case somebody in the future wants to try multiple potential-based layers. It will require more
    # changes than increasing this number.
    self.num_recurrent_layers = 1

    self.input_size = input_size
    self._hidden_size = hidden_size // self.num_recurrent_layers  # hidden size per layer
    self.output_size = self._hidden_size if output_size is None else output_size
    self._stimuli_external = torch.zeros(self.hidden_size, device=device)
    self._stimuli_internal = torch.zeros(self.hidden_size, device=device)

    # Create the common layers.
    self.input_embedding = input_embedding or nn.Linear(self.input_size, self._hidden_size, bias=False)
    self.output_embedding = output_embedding or nn.Linear(self._hidden_size, self.output_size, bias=False)

    # Initialize the values of the potentials.
    if potentials_init is not None:
        self._potentials_init = potentials_init.detach().clone().to(device=device)
    else:
        if activation_nonlin is torch.sigmoid:
            self._potentials_init = -7 * torch.ones(1, self.hidden_size, device=device)
        else:
            self._potentials_init = torch.zeros(1, self.hidden_size, device=device)

    # Initialize the potentials' resting level, i.e., the asymptotic level without stimuli.
    self.resting_level = nn.Parameter(torch.randn(self.hidden_size, device=device))

    # Initialize the potential dynamics' time constant.
    self.tau_learnable = tau_learnable
    if tau_init <= 0:
        raise ValueError("The time constant tau must be initialized positive.")
    self._tau_opt_init = PotentialBased.transform_to_opt_space(
        torch.as_tensor(tau_init - PotentialBased._tau_min, device=device, dtype=torch.get_default_dtype())
    )
    self._tau_opt = nn.Parameter(self._tau_opt_init.reshape(-1), requires_grad=self.tau_learnable)

    # Initialize the potential dynamics' cubic decay.
    self.kappa_learnable = kappa_learnable
    if kappa_init < 0:
        raise ValueError("The cubic decay kappa must be initialized non-negative.")
    self._kappa_opt_init = PotentialBased.transform_to_opt_space(
        torch.as_tensor(kappa_init, device=device, dtype=torch.get_default_dtype()).reshape(-1)
    )
    self._kappa_opt = nn.Parameter(self._kappa_opt_init, requires_grad=self.kappa_learnable)

device: torch.device property

Get the device this model is located on. This assumes that all parts are located on the same device.

hidden_size: int property

Get the number of neurons in the neural field layer, i.e., the ones with the in-/exhibition dynamics.

kappa: Union[torch.Tensor, nn.Parameter] property

Get the cubic decay parameter \(\kappa\).

param_values: torch.Tensor property writable

Get the module's parameters as a 1-dimensional array. The values are copied, thus modifying the return value does not propagate back to the module parameters.

stimuli_external: torch.Tensor property

Get the neurons' external stimuli, resulting from the current inputs. This property is useful for recording during a simulation / rollout.

stimuli_internal: torch.Tensor property

Get the neurons' internal stimuli, resulting from the previous activations of the neurons. This property is useful for recording during a simulation / rollout.

tau: Union[torch.Tensor, nn.Parameter] property

Get the timescale parameter, called \(\tau\) in the original paper [Amari_77].

transform_to_img_space: Callable[[torch.Tensor], torch.Tensor] = torch.exp class-attribute instance-attribute

Function to map parameters to the image space of the original problem.

transform_to_opt_space: Callable[[torch.Tensor], torch.Tensor] = torch.log class-attribute instance-attribute

Function to map parameters to the optimization space.

forward(inputs, hidden=None)

Compute the external and internal stimuli, advance the potential dynamics for one time step, and return the model's output for several time steps in a row.

This method essentially calls forward_one_step several times in a row.

Parameters:

Name Type Description Default
inputs Tensor

Inputs of shape (batch_size, num_steps, dim_input) to evaluate the network on.

required
hidden Optional[Tensor]

Initial values of the hidden states, i.e., the potentials. By default, the network initialized the hidden state to be all zeros. However, via this argument one can set a specific initial value for the potentials. Depending on the shape of inputs, hidden is of shape (hidden_size,) if the input was not batched, else of shape (batch_size, hidden_size).

None

Returns:

Type Description
Tensor

The outputs, i.e., the (linearly combined) activations, and all intermediate potential values, both of

Tensor

shape (batch_size, num_steps, dim_output).

Source code in neuralfields/potential_based.py
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
def forward(self, inputs: torch.Tensor, hidden: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute the external and internal stimuli, advance the potential dynamics for one time step, and return
    the model's output for several time steps in a row.

    This method essentially calls [forward_one_step][neuralfields.PotentialBased.forward_one_step] several times
    in a row.

    Args:
        inputs: Inputs of shape `(batch_size, num_steps, dim_input)` to evaluate the network on.
        hidden: Initial values of the hidden states, i.e., the potentials. By default, the network initialized
            the hidden state to be all zeros. However, via this argument one can set a specific initial value
            for the potentials. Depending on the shape of `inputs`, `hidden` is of shape `(hidden_size,)` if
            the input was not batched, else of shape `(batch_size, hidden_size)`.

    Returns:
        The outputs, i.e., the (linearly combined) activations, and all intermediate potential values, both of
        shape `(batch_size, num_steps, dim_output)`.
    """
    # Bring the sequence of inputs into the shape (batch_size, num_steps, dim_input).
    batch_size = PotentialBased._infer_batch_size(inputs)
    inputs = inputs.view(batch_size, -1, self.input_size)  # moved to the desired device by forward_one_step() later

    # If given use the hidden tensor, i.e., the potentials of the last step, else initialize them.
    hidden = self.init_hidden(batch_size, hidden)  # moved to the desired device by forward_one_step() later

    # Iterate over the time dimension. Do this in parallel for all batched which are still along the 1st dimension.
    inputs = inputs.permute(1, 0, 2)  # move time to first dimension for easy iterating
    outputs_all = []
    hidden_all = []
    for inp in inputs:
        outputs, hidden_next = self.forward_one_step(inp, hidden)
        hidden = hidden_next.clone()
        outputs_all.append(outputs)
        hidden_all.append(hidden_next)

    # Return the outputs and hidden states, both stacked along the time dimension.
    return torch.stack(outputs_all, dim=1), torch.stack(hidden_all, dim=1)

forward_one_step(inputs, hidden=None) abstractmethod

Compute the external and internal stimuli, advance the potential dynamics for one time step, and return the model's output.

Parameters:

Name Type Description Default
inputs Tensor

Inputs of the current time step, of shape (input_size,), or (batch_size, input_size).

required
hidden Optional[Tensor]

Hidden state which are for the model in this package the potentials, of shape (hidden_size,), or (batch_size, input_size). Pass None to leave the initialization to the network which uses init_hidden is called.

None

Returns:

Type Description
Tensor

The outputs, i.e., the (linearly combined) activations, and the most recent potential values, both of shape

Tensor

(batch_size, input_size).

Source code in neuralfields/potential_based.py
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
@abstractmethod
def forward_one_step(
    self, inputs: torch.Tensor, hidden: Optional[torch.Tensor] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute the external and internal stimuli, advance the potential dynamics for one time step, and return
    the model's output.

    Args:
        inputs: Inputs of the current time step, of shape `(input_size,)`, or `(batch_size, input_size)`.
        hidden: Hidden state which are for the model in this package the potentials, of shape `(hidden_size,)`, or
            `(batch_size, input_size)`. Pass `None` to leave the initialization to the network which uses
            [init_hidden][neuralfields.PotentialBased.init_hidden] is called.

    Returns:
        The outputs, i.e., the (linearly combined) activations, and the most recent potential values, both of shape
        `(batch_size, input_size)`.
    """

init_hidden(batch_size=None, potentials_init=None)

Provide initial values for the hidden parameters. This usually is a zero tensor.

Parameters:

Name Type Description Default
batch_size Optional[int]

Number of batches, i.e., states to track in parallel.

None
potentials_init Optional[Tensor]

Initial values for the potentials to override the networks default values with.

None

Returns:

Type Description
Union[Tensor, Parameter]

Tensor of shape (hidden_size,) if hidden was not batched, else of shape (batch_size, hidden_size).

Source code in neuralfields/potential_based.py
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
def init_hidden(
    self, batch_size: Optional[int] = None, potentials_init: Optional[torch.Tensor] = None
) -> Union[torch.Tensor, torch.nn.Parameter]:
    """Provide initial values for the hidden parameters. This usually is a zero tensor.

    Args:
        batch_size: Number of batches, i.e., states to track in parallel.
        potentials_init: Initial values for the potentials to override the networks default values with.

    Returns:
        Tensor of shape `(hidden_size,)` if `hidden` was not batched, else of shape `(batch_size, hidden_size)`.
    """
    if potentials_init is None:
        if batch_size is None:
            return self._potentials_init.view(-1)
        return self._potentials_init.repeat(batch_size, 1).to(device=self.device)

    return potentials_init.to(device=self.device)

potentials_dot(potentials, stimuli) abstractmethod

Compute the derivative of the neurons' potentials w.r.t. time.

Parameters:

Name Type Description Default
potentials Tensor

Potential values at the current point in time, of shape (hidden_size,).

required
stimuli Tensor

Sum of external and internal stimuli at the current point in time, of shape (hidden_size,).

required

Returns:

Type Description
Tensor

Time derivative of the potentials $ rac{dp}{dt}$, of shape (hidden_size,).

Source code in neuralfields/potential_based.py
169
170
171
172
173
174
175
176
177
178
179
@abstractmethod
def potentials_dot(self, potentials: torch.Tensor, stimuli: torch.Tensor) -> torch.Tensor:
    """Compute the derivative of the neurons' potentials w.r.t. time.

    Args:
        potentials: Potential values at the current point in time, of shape `(hidden_size,)`.
        stimuli: Sum of external and internal stimuli at the current point in time, of shape `(hidden_size,)`.

    Returns:
        Time derivative of the potentials $\frac{dp}{dt}$, of shape `(hidden_size,)`.
    """