wenmin-wu/timeseries-multimodal-trajectory-head

Overview

Rather than using K separate prediction heads (one per mode) or a complex MoE router, pack all K trajectory predictions plus K confidence logits into a single linear head with K*T*2 + K outputs. Slice the output tensor into two chunks, reshape the first into (B, K, T, 2), and softmax the second into mode probabilities. This keeps the model architecture simple, parameters shared across modes, and adds one Linear layer to any CNN/Transformer backbone. Pairs directly with the K-mode Gaussian NLL loss.

Quick Start

import torch
import torch.nn as nn
from torchvision.models import resnet34

class MultiModalTrajModel(nn.Module):
    def __init__(self, backbone_features=512, num_modes=3, future_len=50):
        super().__init__()
        self.num_modes = num_modes
        self.future_len = future_len
        self.num_preds = num_modes * 2 * future_len
        self.backbone = resnet34(pretrained=True)
        self.backbone.fc = nn.Identity()
        self.head = nn.Linear(backbone_features, self.num_preds + num_modes)

    def forward(self, x):
        f = self.backbone(x)
        out = self.head(f)                              # (B, K*T*2 + K)
        pred, conf = out[:, :self.num_preds], out[:, self.num_preds:]
        pred = pred.view(-1, self.num_modes, self.future_len, 2)
        conf = torch.softmax(conf, dim=1)
        return pred, conf

Workflow

1. Set the output dim of the final linear layer to num_modes * 2 * future_len + num_modes
2. Split the output on the feature axis at num_modes * 2 * future_len
3. Reshape the first chunk to (B, K, T, 2) — the K candidate trajectories
4. Apply softmax along dim=1 to the second chunk — the mode mixture weights
5. Feed (pred, conf) into the K-mode NLL loss

Key Decisions

• Shared backbone, single head: K modes share all features; the head only has ~K*T*2 extra weights vs. a single-mode model.
• Softmax, not sigmoid: confidences must sum to 1 to form a valid mixture.
• Reshape via .view(): requires contiguous tensors. Use .reshape() if upstream ops produce non-contiguous output.
• vs. K independent heads: K heads produces K independent feature transforms that often collapse to the same mean mode. Shared features + slice-and-reshape forces diversity via the NLL loss instead.

References

• Pytorch Baseline - Train
• Lyft: Complete train and prediction pipeline