4. Implement a Shared-Trunk Multi-Head Regressor

Designing the Shared-Trunk Multi-Head Architecture

A shared-trunk multi-head regressor splits a neural network into two functional zones. The trunk is a stack of layers that every restraint response passes through; it learns a common representation of the input features. The heads are small, independent sub-networks, one per restraint response channel, each mapping the shared representation to a single scalar prediction. The architectural contract is precise: all heads must see the same trunk output, and no head may influence another head's parameters through the trunk during a forward pass. This case builds that structure in PyTorch, then verifies parameter isolation and trunk sharing programmatically rather than by inspection.

The restraint response problem has K output channels (for example, peak force, residual displacement, and energy absorption). Each channel has its own physical scale, so heads should not share final-layer weights. But the input features — geometry, material flags, loading conditions — are common, so a shared trunk is the right inductive bias. The design decision that matters here is where the trunk ends and the heads begin. We place the split after a fixed-width feature block, so every head receives an identical tensor of shape [batch, trunk_dim].

Implementing the Module

The implementation uses nn.ModuleDict for the heads so each head is addressable by name and its parameters are cleanly namespaced. The trunk is a plain nn.Sequential. The forward pass computes the trunk output once, then loops over heads. Because the same tensor z is passed to every head, the trunk is shared by construction; because each head is a distinct module, their parameters are disjoint by construction.

import torch
import torch.nn as nn

class SharedTrunkMultiHeadRegressor(nn.Module):
    def __init__(self, in_features: int, trunk_dim: int = 64,
                 hidden_dim: int = 128, num_outputs: int = 3):
        super().__init__()
        self.num_outputs = num_outputs
        self.trunk = nn.Sequential(
            nn.Linear(in_features, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, trunk_dim),
            nn.ReLU(),
        )
        # One independent head per restraint response channel.
        self.heads = nn.ModuleDict({
            f"head_{k}": nn.Linear(trunk_dim, 1) for k in range(num_outputs)
        })

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = self.trunk(x)                       # [B, trunk_dim], computed once
        outputs = [self.heads[f"head_{k}"](z) for k in range(self.num_outputs)]
        return torch.cat(outputs, dim=1)        # [B, num_outputs]

model = SharedTrunkMultiHeadRegressor(in_features=12, num_outputs=3)
x = torch.randn(8, 12)
y_hat = model(x)
print(y_hat.shape)  # torch.Size([8, 3])

Two details are worth calling out. First, torch.cat(outputs, dim=1) produces a single tensor of shape [B, K], which is exactly the layout the masked-MSE loss in the next case expects. Second, the trunk is invoked exactly once per forward pass — if you accidentally call self.trunk(x) inside the head loop, you would recompute the trunk K times, wasting compute and, more subtly, breaking the guarantee that all heads consume the identical representation.

Verifying Parameter Isolation and Trunk Sharing

Architecture claims should be tested, not assumed. The verification below checks three properties: (1) the trunk parameters are shared across all heads, (2) each head's parameters are disjoint from every other head's, and (3) a gradient flowing into one head's loss does not update another head's weights.

def parameter_report(model):
    trunk_params = {id(p) for p in model.trunk.parameters()}
    head_params = {
        name: {id(p) for p in head.parameters()}
        for name, head in model.heads.items()
    }
    # 1. Heads must not overlap with the trunk.
    for name, ids in head_params.items():
        assert trunk_params.isdisjoint(ids), f"{name} shares params with trunk"
    # 2. Heads must be pairwise disjoint.
    names = list(head_params)
    for i in range(len(names)):
        for j in range(i + 1, len(names)):
            assert head_params[names[i]].isdisjoint(head_params[names[j]]), \
                f"{names[i]} and {names[j]} share parameters"
    print("isolation OK:", len(trunk_params), "trunk params,",
          sum(len(v) for v in head_params.values()), "head params")

parameter_report(model)

# 3. Gradient isolation: train only head_0, confirm head_1/head_2 grads stay None.
model.zero_grad()
loss = model.heads["head_0"](model.trunk(x)).pow(2).mean()
loss.backward()
for name, head in model.heads.items():
    grads = [p.grad is not None for p in head.parameters()]
    print(name, "grads present:", grads)

Expected output from the gradient check: head_0 grads present: [True, True], while head_1 and head_2 report [False, False]. This is the concrete proof that head-specific parameters are isolated. Note that the trunk will receive gradients in this test, because head_0 backpropagates through it — that is the intended sharing behavior, and it is why the trunk is described as shared rather than frozen.

Common Failure Symptoms

Checklist for This Build

• Trunk output shape is [B, trunk_dim] and is computed exactly once per forward pass.
• Each restraint response has its own nn.Linear(trunk_dim, 1) head under a distinct ModuleDict key.
• parameter_report passes: trunk and heads are disjoint, heads are pairwise disjoint.
• Gradient isolation test shows non-target heads with grad is None.
• Forward output is [B, K], ready for the masked-MSE loss in the next stage.

With the module built and its isolation properties verified, the architecture is ready to be wired into the masked-loss training loop. The key invariant to carry forward is that the trunk is the only shared parameter path; any future change that couples heads must be deliberate and documented.