def __init__(self, q=1, d=2, m=3, u=4, delta=3,
gamma_alpha=0.1, tau_alpha=0.5, tau_mu=100, tau_eta=100,
upd_m="WTA", upd_k="ad_hoc_WTA",
beta_k=0.001,
psi_fn="identity",
scramble=False):
"""
:param q: number of neurons
:param d: size of each key
:param m: number of keys/memory units
:param u: size of each memory unit
:param gamma_alpha: softmax temperature (key matching)
:param tau_alpha: threshold on the attention score of the winning key, to eventually trigger scrambling
:param tau_mu: number of steps below which a key is considered to be not-used enough
:param tau_eta: number of steps after which a key is considered old
:param delta: number of top attention responses to select (top-delta)
:param upd_m: update memory strategy (None, 'WTA')
:param upd_k: update key strategy (None, 'ad_hoc_WTA', 'grad_WTA')
:param beta_k: learning rate for key-update purposes when upd_k is 'ad_hoc_WTA'
:param psi_fn: function to project the neuron input onto the key space
:param scramble: triggers the key/memory scrambling routine when upd_k is 'ad_hoc_WTA'
"""
super(CNUs, self).__init__()
assert upd_m in (None, 'WTA'), "Unknown value for upd_m, it must be None or 'WTA'"
assert upd_k in (None, 'ad_hoc_WTA', 'grad_WTA'), "Unknown value for upd_k, it must be " \
"None, 'ad_hoc_WTA', or 'grad_WTA'"
assert upd_m is None or (upd_m == 'WTA' and upd_k is not None), \
"If upd_m is 'WTA', then upd_k must be ad_hoc_WTA or grad_WTA (it cannot be None)"
self.q = q
self.d = d
self.m = m
self.u = u
self.gamma_alpha = gamma_alpha
self.tau_alpha = tau_alpha
self.tau_mu = tau_mu
self.tau_eta = tau_eta
self.scramble = scramble
self.delta = min(delta, self.m)
self.upd_m = upd_m
self.upd_k = upd_k
self.beta_k = beta_k
self.psi_fn = psi_fn
self.debug = False # Temporarily used
self.reset_memories = True
# Creating keys (self.K) and memories (self.M)
self.M = torch.nn.Parameter(torch.empty((self.q, self.m, self.u), dtype=torch.float32))
if self.upd_k == "ad_hoc_WTA":
self.register_buffer('K', torch.zeros((self.q, self.m, self.d)))
else:
self.K = torch.nn.Parameter(torch.empty((self.q, self.m, self.d), dtype=torch.float32))
# Buffers for ad_hoc_WTA key updates (average usefulness register buffer "mu" and age "eta")
if self.upd_k == "ad_hoc_WTA":
self.register_buffer('mu', torch.zeros(self.q, m, dtype=torch.float))
self.register_buffer('eta', torch.ones((self.q, m), dtype=torch.float) * self.tau_eta)
if self.debug:
self.register_buffer('key_counter', torch.zeros(self.q, m, dtype=torch.float))
else:
self.mu = None
self.eta = None
# Scrambling stats
self.register_buffer('scrambling_count', torch.zeros(self.q, dtype=torch.long))
# Initializing memories and keys
self.reset_parameters()