Architecture Methods › Code Implementation: HAT Masks
Language AI HandbookAll chapters Architecture Methods › Code Implementation: HAT MasksOne section · Sections in this chapter8 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Now let us implement the mask-based protection mechanism from HAT. This gives us a fixed-capacity network where each task learns a sparse attention mask over units, and the accumulated masks protect previously claimed capacity from future modification. The implementation has three key components: the task embedding that parameterizes the mask, the annealing schedule that hardens the mask over training, and the gradient zeroing that enforces protection at each update step.
class HATNetwork(nn.Module):
"""
Simplified HAT implementation: a 2-layer MLP with per-task sigmoid attention masks.
Task embeddings control which units are active for each task.
"""
def __init__(self, input_dim, hidden_dim, output_dim, max_tasks=10):
super().__init__()
self.hidden_dim = hidden_dim
# Shared backbone weights
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, output_dim)
self.relu = nn.ReLU()
# Task embeddings: one scalar per hidden unit per task
self.te1 = nn.Embedding(max_tasks, hidden_dim)
nn.init.normal_(self.te1.weight, mean=0.0, std=0.01)
# Cumulative mask: union of all prior tasks' hard masks
self.register_buffer("cum_mask1", torch.zeros(hidden_dim))
def get_mask(self, task_id, s=1.0):
"""Compute soft attention mask for given task and temperature s."""
tid = torch.tensor(task_id, dtype=torch.long)
a1 = torch.sigmoid(s * self.te1(tid))
return a1
def forward(self, x, task_id, s=1.0):
a1 = self.get_mask(task_id, s)
h1 = self.relu(self.fc1(x) * a1)
return self.fc2(h1), a1
def compute_hat_loss(logits, targets, a1, cum_mask1, lam=0.75):
"""
HAT loss = cross-entropy + sparsity penalty on unclaimed units.
We encourage the new task's mask to be sparse in the unclaimed region.
"""
ce_loss = nn.functional.cross_entropy(logits, targets)
unclaimed = 1.0 - cum_mask1.detach()
sparsity = lam * (a1 * unclaimed).sum() / a1.numel()
return ce_loss + sparsityThe task embedding te1 stores one scalar per hidden unit per task. When the temperature is small, the sigmoid output is soft (near 0.5 for all units). As grows during training, the output converges toward 0 or 1, producing a hard binary mask at convergence. The small normal initialization (std=0.01) places all initial embeddings close to zero, meaning the initial soft mask values are near . All units start as roughly equally available before training shapes the mask.
def train_hat_task(model, task_id, X_tr, y_tr, epochs=40, lr=5e-3, s_max=100.0):
"""Train HAT model on a new task with mask annealing."""
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
n_epochs = epochs
for epoch in range(n_epochs):
# Anneal temperature from 1 to s_max over training
s = 1.0 + (s_max - 1.0) * epoch / (n_epochs - 1)
model.train()
optimizer.zero_grad()
logits, a1 = model(X_tr, task_id, s=s)
loss = compute_hat_loss(logits, y_tr, a1, model.cum_mask1)
loss.backward()
# Gradient protection: zero gradients for units claimed by prior tasks
with torch.no_grad():
if model.fc1.weight.grad is not None:
for neuron_idx in range(model.hidden_dim):
if model.cum_mask1[neuron_idx] > 0.5:
model.fc1.weight.grad[neuron_idx, :] *= 0.0
if model.fc1.bias.grad is not None:
model.fc1.bias.grad[neuron_idx] *= 0.0
optimizer.step()
# Binarize the final mask and update cumulative mask
with torch.no_grad():
hard_mask = (model.get_mask(task_id, s=s_max) > 0.5).float()
model.cum_mask1 = torch.clamp(model.cum_mask1 + hard_mask, 0.0, 1.0)
def evaluate_hat(model, task_id, X_te, y_te, s=100.0):
"""Evaluate HAT model on a specific task."""
model.eval()
with torch.no_grad():
logits, _ = model(X_te, task_id, s=s)
preds = logits.argmax(dim=1)
acc = (preds == y_te).float().mean().item()
return accThe gradient protection loop is the core of the HAT mechanism. After backpropagation computes gradients for all parameters, we zero out the row of fc1.weight.grad corresponding to any neuron that has been claimed by a prior task (cumulative mask value above 0.5). This prevents the optimizer from updating those neurons' incoming weights, effectively freezing their receptive fields. The bias gradient is also zeroed to maintain consistency.
hat_net = HATNetwork(input_dim=10, hidden_dim=64, output_dim=2)
# Train on Task 1
train_hat_task(hat_net, task_id=0, X_tr=X1_tr, y_tr=y1_tr, epochs=50)
hat_t1_before = evaluate_hat(hat_net, task_id=0, X_te=X1_te, y_te=y1_te)
claimed_after_t1 = hat_net.cum_mask1.sum().item()
# Train on Task 2
train_hat_task(hat_net, task_id=1, X_tr=X2_tr, y_tr=y2_tr, epochs=50)
hat_t1_after = evaluate_hat(hat_net, task_id=0, X_te=X1_te, y_te=y1_te)
hat_t2 = evaluate_hat(hat_net, task_id=1, X_te=X2_te, y_te=y2_te)
claimed_after_t2 = hat_net.cum_mask1.sum().item()HAT Network Results Task 1 accuracy before Task 2 training: 0.7812 Task 1 accuracy after Task 2 training: 0.8438 Task 2 accuracy: 0.9250 Units claimed after Task 1: 7/64 Units claimed after Task 2: 13/64 Remaining unclaimed units: 51/64
The cumulative mask grows as more tasks are added. Each task carves out a portion of the hidden layer's units, and those units are subsequently protected from modification. The remaining unclaimed units are available for future tasks. When the unclaimed region shrinks to a very small number of units, the network will struggle to accommodate new tasks effectively.
Notice that unlike the progressive network, HAT's protection is not perfect in this simplified implementation. In the full HAT paper, gradient protection is applied to all weight tensors connected to protected neurons, including the output layer and all lateral connections. Our simplified version focuses on fc1 to illustrate the principle.