Multi-input and multi-output models

Meera’s new question was really two questions about the same thing. Every evening, each of the six stores plans tomorrow’s bake for twelve pastries. For every pastry, the planner wants to know: how many will we sell? And: will we sell out before 6pm? Selling out of croissants at 4pm is fine. Selling out of the pistachio croissants, Crumb’s new bestseller, at 11am means a queue of disappointed customers and an Instagram story about it.

“Two questions, so two models,” Arjun said. He had last lesson’s forecaster. He copied it, changed the last layer and the loss, and had a sell-out classifier by lunchtime.

By teatime, the two models were contradicting each other. For Saturday’s pistachio croissants at the Baner store, the plan was to bake 26. The sales model said “you’ll sell about 18”. The sell-out model said “79% chance you’ll sell out”. You cannot sell 18 out of 26 and sell out.

“They each learned their own idea of what Saturday at Baner looks like,” Priya said. “Two students, two textbooks, two different answers. Let’s give them one textbook.”

Why one model with two heads

Priya’s analogy: Crumb’s central kitchen preps once for many dishes. The same tray of laminated dough becomes croissants, pain au chocolat and danishes. Nobody laminates three separate batches. The expensive, shared work happens once; each dish adds its own finishing touch.

A multi-task model works the same way. The expensive part, understanding “what is going on with pistachio croissants at Baner this week”, is learned once, in a shared trunk. Each question gets a small output head (a final layer or two that turns the shared representation into one particular answer). Three benefits:

• Consistency. Both answers come from the same understanding, so they cannot drift into contradiction as easily.
• Better learning. Each task is extra training signal for the shared part. Knowing which days sold out tells the model that sales on those days were capped, which helps it understand sales.
• One thing to run. One model to deploy, monitor and retrain instead of two.

And the inputs are not all the same kind of data either. Meera’s planners know three different kinds of things about each pastry-store-day:

Input	Kind	Example for pistachio croissant, Baner, Saturday
Last 14 days of sales, sell-outs and rain	A sequence	[20, 40, 18, 18, 9, 23, 16, ...]
Which pastry, which store	Categories (IDs)	product 3, store 1
Price, how long since launch, tomorrow’s weekday, rain forecast, planned bake	Plain numbers	₹180, 316 days, Saturday, dry, 26

Each kind needs a different reader. A sequence needs something that reads in order, like the GRU from two lessons ago. Categories need a lookup. Plain numbers need a small dense network. This is a multi-input model: separate encoders for each input, combined in the middle.

Turning categories into numbers: embeddings

“Pastry number 3” is not a quantity. Pistachio croissant (3) is not three times a croissant (1). If you feed the ID as a number, the model will assume it is.

The fix is an embedding: a small learned vector for each category, stored in a lookup table. nn.Embedding(12, 4) holds 12 rows of 4 numbers, one row per pastry, starting random and learned during training like any other weights. Asking for pastry 3 returns row 3. After training, pastries that behave alike (sell out on weekends, collapse in the rain) end up with similar vectors, exactly like word embeddings in Module 07, but learned for this task.

from torch import nn
import torch

product_emb = nn.Embedding(num_embeddings=12, embedding_dim=4)   # 12 pastries -> 4 numbers each
print(product_emb(torch.tensor([3])))     # the 4 numbers for the pistachio croissant
print(product_emb.weight.shape)           # torch.Size([12, 4]): 48 learnable parameters

The model, box by box

flowchart LR
    A["14-day history"] --> B["GRU encoder"]
    C["Product ID"] --> D["Product embedding"]
    E["Store ID"] --> F["Store embedding"]
    G["Price, age, weekday, rain, planned bake"] --> H["Metadata MLP"]
    D --> H
    F --> H
    B --> I["Fuse: concatenate"]
    H --> I
    I --> J["Shared trunk"]
    J --> K["Units head"]
    J --> L["Sell-out head"]

Combining the encoders’ outputs is called fusion. The simplest kind, used here, is concatenation: put the GRU’s 32 numbers and the metadata encoder’s 32 numbers side by side into one vector of 64. Because it happens after each input has been encoded separately, this is sometimes called late fusion (as opposed to early fusion, which glues raw inputs together before any encoding). Other fusion methods exist (adding, multiplying, letting one input attend to another), but concatenation is a strong, simple default.

class PastryModel(nn.Module):
    def __init__(self, n_products: int = 12, n_stores: int = 6, hidden: int = 32):
        super().__init__()
        # Input 1: the sequence
        self.seq_encoder = nn.GRU(input_size=3, hidden_size=hidden, batch_first=True)
        # Inputs 2 and 3: categories and plain numbers
        self.product_emb = nn.Embedding(n_products, 4)
        self.store_emb = nn.Embedding(n_stores, 3)
        self.meta_encoder = nn.Sequential(nn.Linear(4 + 3 + 6, 32), nn.ReLU())
        # Shared trunk after fusion
        self.trunk = nn.Sequential(nn.Linear(hidden + 32, 64), nn.ReLU())
        # Two heads
        self.units_head = nn.Linear(64, 1)        # how many will sell (regression)
        self.sellout_head = nn.Linear(64, 1)      # will it sell out (a logit)

    def forward(self, seq, product, store, meta):
        _, h = self.seq_encoder(seq)                                  # h: (1, batch, 32)
        m = self.meta_encoder(torch.cat([self.product_emb(product),
                                         self.store_emb(store), meta], dim=1))   # (batch, 32)
        z = self.trunk(torch.cat([h[-1], m], dim=1))                  # fuse -> (batch, 64)
        return {"units": self.units_head(z), "sellout": self.sellout_head(z)}

Three things are new compared with every model so far. forward takes four arguments instead of one. It returns a dictionary of outputs, so later code says out["sellout"] rather than remembering that the sell-out score is “the second thing”. And the heads are deliberately tiny: almost all the learning happens in the shared parts. The whole model has 8,356 parameters.

Arjun adds the two losses together

Training a two-headed model needs one number to call .backward() on. The obvious move: compute each head’s loss and add them.

• Units: mean squared error between predicted and actual units sold.
• Sell-out: binary cross-entropy (with logits) between the score and whether it sold out.

loss = mse(out["units"], y_units) + bce(out["sellout"], y_sold_out)

Arjun trained it on the first fourteen months of data, validated on the next three and tested on the last three, split by date as in the forecasting lesson. The sales head was decent. The sell-out head was a disaster:

Model	Units MAE	Sell-out AUROC	Sell-out F1
Seasonal naive (sales only)	5.27	n/a	n/a
Two-headed, raw losses added	4.14	0.63	0.00

An F1 of zero means it predicted “won’t sell out” for every single pastry, every day, even though about a third did. AUROC (the chance that a randomly chosen sell-out day gets a higher score than a randomly chosen normal day; 0.5 is a coin flip, 1.0 is perfect) was 0.63, barely better than guessing. The sell-out head had effectively not been trained.

Arjun printed both losses during training:

epoch  1  units loss 135.248  sellout loss 0.615
epoch  5  units loss  42.642  sellout loss 0.603
epoch 15  units loss  26.668  sellout loss 0.606

The units loss fell from 135 to 27. The sell-out loss did not move at all.

Following one pastry through both losses

Priya traced a single batch at the very start of training to put a number on it. With raw units, the starting losses were:

$$L_{\text{units}}=778.8L_{\text{sellout}}=0.71$$

What matters is not the losses but their gradients on the shared trunk, because that is what moves the shared weights. She measured the size of the gradient each loss sends into the trunk’s first layer:

Units target	Gradient from units loss	Gradient from sell-out loss	Ratio
Raw (pastries)	24.73	0.084	296 to 1
Standardised (z-scores)	0.47	0.084	5.7 to 1

With raw units, every update to the shared trunk is 99.7% “help the units head”. The sell-out loss is a whisper in a hurricane. Standardise the units target, the same $z=(x-\mu )/\sigma$ from Module 27 with the training-period mean of 25.7 pastries and standard deviation of 16.6, and the ratio drops from 296 to under 6. Now both voices are heard.

Loss weighting, properly

The combined loss is a weighted sum:

$$L=w_{\text{units}}{L}_{\text{units}}+w_{\text{sellout}}{L}_{\text{sellout}}$$

Choosing those weights is loss weighting, and it is the central design decision in any multi-output model. The practical recipe, in order:

1. Put every loss on a sensible scale first. Standardise regression targets. Classification losses with cross-entropy usually start around 0.7 (two classes) or $\ln (\text{classes})$ (more). Aim for every loss to start somewhere between about 0.1 and 10.
2. Start with equal weights (1 and 1) and watch each loss separately during training, not just the total.
3. Adjust weights to match what matters. If the sell-out question is the one Meera cares most about, give it more weight. Tune on validation metrics for each head, not on the combined loss.
4. Re-check when anything changes. A new head, a new target scale or a new dataset can shift the balance.

There are cleverer schemes. One learns the weights during training from each task’s uncertainty (Kendall, Gal and Cipolla, 2018); another, GradNorm, adjusts weights so each task’s gradients stay balanced. They help when you have many heads. For two or three, steps 1 to 3 are usually enough.

Here is what happened when Arjun standardised the target and tried a few weights:

Model	Units MAE	Sell-out AUROC	Sell-out F1
Seasonal naive	5.27	n/a	n/a
Units-only model	3.85	(untrained)	(untrained)
Sell-out-only model	(untrained)	0.867	0.683
Two heads, raw units, 1 : 1	4.14	0.63	0.00
Two heads, standardised, 1 : 1	3.78	0.868	0.682
Two heads, standardised, 1 : 0.3	3.77	0.867	0.689
Two heads, standardised, 1 : 3	3.81	0.867	0.682

Read the middle rows carefully. Once the scales are fixed, one shared model matches the sell-out specialist and slightly beats the units specialist (3.78 against 3.85). The extra task helped the main one. And once the scales were sensible, the exact weights barely mattered: 0.3, 1 and 3 all land in the same place. Scale is the big lever; weights are fine-tuning.

Following one pastry through the model

Now trace a single prediction from the trained model. Pistachio croissant, Baner store, a dry Saturday in the test period.

Sequence input. The previous seven days’ sales were [20, 40, 18, 18, 9, 23, 16], and the pastry sold out on three of those days. Each day becomes three numbers: standardised sales, a sold-out flag and a rain flag. (The model sees 14 days; these are the last 7.)

Category inputs. Product 3 and store 1 are looked up in their embedding tables: four numbers for the pistachio croissant, three for Baner.

Plain numbers. Price ₹180 (scaled to 2.0), about 316 days since launch (scaled to 0.96), Saturday as sine and cosine (−0.97, −0.22), no rain, and tomorrow’s planned bake of 26 (standardised to 0.02).

Encoders and fusion. The GRU reads the 14 days into 32 numbers. The metadata encoder turns the 13 category and plain numbers into 32 more. Concatenated: 64. The trunk mixes them into a new 64.

Heads. Units: 24.1 croissants. Sell-out: 58%.

What happened. They sold 21 and did not sell out.

Now the part that makes this useful for planning. The planned bake is an input, so the planner can ask “what if?”:

If Baner bakes	Predicted sales	Chance of selling out
26	24.1	58%
30	26.0	32%
40	27.6	5%
50	27.4	1%

The two heads now tell one coherent story: demand is around 26 to 28, so baking 30 keeps the sell-out risk to about one in three, and baking 40 almost guarantees leftovers. Meera’s planner can pick the trade-off, remembering from the forecasting lesson that a missed sale costs about twice a leftover.

The data plumbing

Multi-input models need datasets that return several inputs. The cleanest pattern is to return a dictionary, which PyTorch’s DataLoader batches automatically, key by key:

from torch.utils.data import Dataset

class PastryDays(Dataset):
    def __init__(self, seq, product, store, meta, units, sold_out):
        self.seq, self.product, self.store, self.meta = seq, product, store, meta
        self.units, self.sold_out = units, sold_out

    def __len__(self):
        return len(self.seq)

    def __getitem__(self, i):
        return {
            "seq": self.seq[i],             # (14, 3)
            "product": self.product[i],     # scalar ID
            "store": self.store[i],         # scalar ID
            "meta": self.meta[i],           # (6,)
            "units": self.units[i],         # (1,) standardised
            "sold_out": self.sold_out[i],   # (1,) 0 or 1
        }

And the training step reads by name, computes each loss, logs them separately and combines them with explicit weights:

weights = {"units": 1.0, "sellout": 1.0}
mse, bce = nn.MSELoss(), nn.BCEWithLogitsLoss()

for batch in loader:
    out = model(batch["seq"], batch["product"], batch["store"], batch["meta"])
    losses = {
        "units": mse(out["units"], batch["units"]),
        "sellout": bce(out["sellout"], batch["sold_out"]),
    }
    loss = sum(weights[k] * losses[k] for k in losses)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    # log every component, not just the total
    log = {k: round(v.item(), 3) for k, v in losses.items()}

The sequence, metadata and labels are built from the raw sales with the same sliding-window code as the forecasting lesson, split by date, with the units mean and standard deviation computed from the training period only. To evaluate, convert the units prediction back with pred * std + mean, and score the sell-out head with torchmetrics’ BinaryAUROC and BinaryF1Score.

Where multi-head models go wrong

One loss quietly takes over. You have seen the main case. The symptom is always the same: one head’s loss falls and the other’s is flat. Log each loss separately, every epoch, from the start.

Some examples have only some labels. Next month, Meera adds a third head for customer ratings, but only 10% of days have ratings. Do not drop the other 90% of days. Compute the rating loss only where a label exists, using a mask (multiply the per-example loss by 1 where there is a label and 0 where there is not, then average over the labelled ones). Every multi-head project meets this eventually.

Picking “the best epoch” is now ambiguous. The units head might peak at epoch 12 and the sell-out head at epoch 20. Decide in advance which validation metric chooses the checkpoint (usually the one the business cares about), or use a weighted combination, and write it down.

Tasks that fight. If two heads want genuinely different things from the shared trunk, adding the second makes the first worse. Always compare against single-task models, as Arjun did. If the shared model loses on the main task, give each head more private layers, or split the models.

The capped-sales trap. Units “sold” are not units “wanted” on sell-out days. A model trained to predict sales will learn to predict the cap. The sell-out head helped here, but if you care about true demand, treat sell-out days specially, for example by training with a loss that only penalises under-predictions on those days. This is a well-known problem in retail forecasting called censored demand.

Try it yourself

1. Remove the planned-bake feature from the metadata (set it to zero everywhere) and retrain. What happens to the sell-out head’s AUROC? Why?
2. Train with standardised units but weights of 1 : 0.01. Then 1 : 100. Log both losses. At what point does each head stop learning?
3. Add a third head: “will it sell fewer than half of what was baked?” (a waste warning). Build the label, add a head and a loss with weight 1, and compare all three heads’ metrics against the two-head model.

Show the answer

1. In Arjun’s run, the sell-out AUROC fell from 0.868 to 0.741 and F1 from 0.68 to 0.49 (the units MAE also worsened, from 3.78 to 4.30). Whether a pastry sells out depends on demand and on how many were baked. Without the planned bake, the model can only guess whether tomorrow’s typical demand will exceed a bake quantity it cannot see.
2. At 1 : 0.01, the sell-out loss barely moves and its F1 falls towards zero, as with raw units. At 1 : 100, the units head suffers instead: its MAE rises towards the seasonal naive level. Somewhere in between, both heads learn. That range is usually wide once targets are scaled, as the 0.3 to 3 results showed.
3. The waste head is closely related to the other two, so it should learn well and may help a little. If any existing head gets noticeably worse, you are seeing the trunk’s capacity being shared; try hidden=64.

Common mistakes

• Adding raw losses on different scales. Standardise regression targets and check every loss starts in a comparable range.
• Logging only the total loss. A falling total can hide a head that is not learning at all.
• Feeding category IDs as numbers. Use nn.Embedding for anything that is a label rather than a quantity.
• Fusing raw inputs of very different kinds early. Give each kind of input its own encoder, then fuse.
• Never comparing with single-task models. That comparison is how you know the shared model is helping, not hurting.

Recap

• Multi-input models give each kind of input (sequence, categories, plain numbers, images) its own encoder, then fuse the results, most simply by concatenation.
• Embeddings turn category IDs into small learned vectors.
• Multi-output models share a trunk and put small heads on top; related tasks can improve each other.
• The total loss is a weighted sum; put every loss on a comparable scale first, then tune weights on per-head validation metrics.
• Log each loss separately, handle missing labels with masks, and always compare against single-task baselines.

Arjun walked Meera through the what-if table for the pistachio croissants. “It’s one model that understands each pastry at each store, and it answers both questions from the same understanding, so they agree. Our first try only learned the sales question, because its mistakes were measured in hundreds and the sell-out mistakes in fractions. Once we put them on the same scale, it learned both. And because you tell it how many you plan to bake, it tells you how likely you are to run out.”

“Thirty, then,” said Meera. “For Saturday at Baner.”

Priya was already looking at a pile of crumpled supplier invoices on the next desk. “Next: a model with an image as one input, metadata as another, and a head that reads text. Your project for the rest of the month is to read those.”

Go deeper

• PyTorch tutorials for nn.Embedding, custom datasets and models with several inputs.
• torchmetrics documentation for BinaryAUROC, BinaryF1Score and regression metrics such as MeanAbsoluteError.