Multi-Task Learning for Simultaneous Price, Volume, and Volatility Prediction

Multi-task learning (MTL) is usually sold with an assertion: share an encoder across correlated targets and the primary task gets better. In trading the correlated targets are obvious — returns, volume, and realized volatility all fall out of the same order flow — and the assertion is almost never tested. The interesting question is not whether the tasks are related. It is whether the shared gradients agree, and what happens on the folds where they do not.

This article puts two things at the center that most MTL write-ups treat as footnotes:

1. Loss balancing is the experiment, not a detail. Fixed weights, Kendall uncertainty weighting, and GradNorm are three different models. Run all three on the same folds and report the learned weights alongside the primary-task metric for each.
2. Negative transfer is measurable before you see the metric. The cosine similarity between task gradients on the shared encoder tells you, during training, whether the auxiliary tasks are pulling the representation somewhere the primary task wants to go. Sign the cosines, then check whether the sign predicted the outcome on that fold.

Everything else in the pipeline — the volatility process, the training loop, the leakage controls, the validation protocol — is already covered elsewhere on this blog and is linked rather than re-derived.

Setup

Given input features $\mathbf{x} \in \mathbb{R}^d$ (OHLCV, technical indicators, order flow), three targets:

• Task 1 (primary): next-period return $y^{(1)} = r_{t+1}$
• Task 2 (auxiliary): next-period log volume $y^{(2)} = \log V_{t+1}$
• Task 3 (auxiliary): next-period realized volatility $y^{(3)} = \sigma_{t+1}$

A multi-task model produces all three simultaneously, $\hat{\mathbf{y}} = f_\theta(\mathbf{x})$, and the multi-task risk is a weighted sum of per-task risks:

$\mathcal{R}_{\text{MTL}}(\theta) = \sum_{k=1}^{K} w_k \cdot \mathbb{E}\bigl[\ell^{(k)}(f_\theta^{(k)}(\mathbf{x}), y^{(k)})\bigr]$

The whole article is about the $w_k$ and about what the per-task gradients do to each other.

Why joint training might help, in one paragraph. Auxiliary tasks constrain the shared representation to explain more than one market phenomenon, which is a capacity control and an inductive bias at the same time; and because volume and volatility are directly observed while "expected return" is not, the auxiliary heads supply cleaner gradient signal than the primary head does. The case for one model emitting many outputs is argued at length — with the interpretability machinery attached — in temporal fusion transformers for multi-horizon forecasting, which makes the same shared-encoder-many-heads argument for multi-horizon quantiles.

Architecture, briefly

Hard parameter sharing: a shared encoder $g_\phi$ feeds $K$ task-specific heads $h_{\psi_k}$, so $\hat{y}^{(k)} = h_{\psi_k}(g_\phi(\mathbf{x}))$. This is the version measured here, because it is the version where gradient conflict on $\phi$ is well-defined.

Soft parameter sharing gives each task its own encoder with a coupling penalty $\lambda \sum_{k \neq j} \|\phi_k - \phi_j\|^2$ — more parameters, more flexibility, and no single shared parameter vector to measure conflict on. Cross-stitch networks sit in between, mixing per-task features through a learned matrix $\alpha_{kj}$ at each level. Both are worth trying if hard sharing shows conflict, and both are outside the scope of the measurement below.

The Experiment That Matters: Three Loss-Balancing Schemes

The naive loss $\mathcal{L} = \sum_k w_k \mathcal{L}^{(k)}$ is scale-sensitive. If return loss lives around $0.01$ and volume loss around $1.0$, volume owns the gradient and the return head starves. Three responses:

Fixed weights. Set $w_k = 1$ after standardizing every target. The honest baseline — if it wins, the adaptive schemes are ceremony.

Uncertainty weighting (Kendall et al., 2018). Learn a homoscedastic noise scale $\sigma_k$ per task:

$\mathcal{L}_{\text{MTL}} = \sum_{k=1}^{K} \frac{1}{2\sigma_k^2} \mathcal{L}^{(k)} + \log \sigma_k$

High-uncertainty tasks get down-weighted automatically; the $\log \sigma_k$ term blocks the trivial $\sigma_k \to \infty$ solution. Note this $\sigma_k$ is a training-time loss-weighting device, not a predictive interval — for uncertainty you can actually size a position with, see conformal prediction.

GradNorm (Chen et al., 2018). Balance gradient magnitudes rather than loss scales. Each step: compute $G_k = \|\nabla_\phi w_k \mathcal{L}^{(k)}\|_2$ and the mean $\bar{G}$, compute the relative training rate $\tilde{r}_k = \frac{\mathcal{L}^{(k)}(t)/\mathcal{L}^{(k)}(0)}{\bar{r}}$, and update $w_k \leftarrow w_k - \eta_w \nabla_{w_k} \sum_k |G_k - \bar{G} \cdot \tilde{r}_k^\alpha|$. All tasks then train at comparable rates regardless of loss scale.

The MTL-specific code is the heads, the list-returning forward, and the loss aggregation. The Linear/BatchNorm/ReLU/Dropout stack, the Adam/cosine/clip boilerplate, and the epoch loop are the standard pattern shown in DeepLOB and are omitted here.

import torch
import torch.nn as nn


class MultiTaskTradingModel(nn.Module):
    """Hard parameter sharing: one encoder, K heads."""

    def __init__(self, encoder: nn.Module, repr_dim: int, n_tasks: int = 3):
        super().__init__()
        self.shared_encoder = encoder          # any MLP/CNN/GRU trunk
        self.task_heads = nn.ModuleList(
            nn.Linear(repr_dim, 1) for _ in range(n_tasks)
        )

    def forward(self, x):
        h = self.shared_encoder(x)
        return [head(h).squeeze(-1) for head in self.task_heads]

    def shared_repr(self, x):
        return self.shared_encoder(x)


class UncertaintyWeightedLoss(nn.Module):
    """Kendall et al. (2018) homoscedastic weighting."""

    def __init__(self, n_tasks: int = 3):
        super().__init__()
        self.log_vars = nn.Parameter(torch.zeros(n_tasks))  # log(sigma^2)

    def forward(self, losses: list) -> torch.Tensor:
        return sum(
            torch.exp(-self.log_vars[i]) * loss + self.log_vars[i]
            for i, loss in enumerate(losses)
        )

    def get_weights(self) -> list:
        with torch.no_grad():
            return [torch.exp(-lv).item() for lv in self.log_vars]

UncertaintyWeightedLoss has parameters, so it must go into the optimizer alongside the model: optim.Adam(list(model.parameters()) + list(uw.parameters()), ...). Forgetting this is the most common way to "run uncertainty weighting" and silently run fixed weights instead.

What to report

For each scheme, on each fold: the learned final task weights, the primary-task metric, and — because a weighting scheme is a model choice — how many schemes were compared before picking one.

Three schemes times several folds is already a small model search. Any improvement reported here has to survive the multiple-testing correction described in deflated Sharpe and multiple testing before it means anything.

Negative Transfer: Sign the Gradients

This is the part worth keeping. Negative transfer is when auxiliary tasks make the primary task worse, and it has a direct diagnostic: the angle between task gradients in the shared parameter space.

$\cos\bigl(\nabla_\phi \mathcal{L}^{(k)}, \nabla_\phi \mathcal{L}^{(j)}\bigr) < 0 \implies \text{conflicting tasks}$

Measured on the shared encoder only — the heads are task-specific by construction and always "agree" trivially.

import torch.nn.functional as F


def shared_grad(model, x, y, task_idx, criterion=nn.MSELoss()):
    """Gradient of task `task_idx` w.r.t. the shared encoder, flattened."""
    model.zero_grad(set_to_none=True)
    loss = criterion(model(x)[task_idx], y)
    loss.backward()
    return torch.cat([
        p.grad.detach().flatten()
        for p in model.shared_encoder.parameters()
        if p.grad is not None
    ])


def task_conflict(model, x, y_by_task, task_names):
    """Pairwise cosine similarity between per-task shared-encoder gradients."""
    grads = {
        name: shared_grad(model, x, y_by_task[name], i)
        for i, name in enumerate(task_names)
    }
    return {
        (a, b): F.cosine_similarity(
            grads[a].unsqueeze(0), grads[b].unsqueeze(0)
        ).item()
        for i, a in enumerate(task_names)
        for b in task_names[i + 1:]
    }

Call this on a held-out batch at a fixed cadence during training, not once at the end. A pair can start aligned and diverge as the encoder specializes; a single end-of-training number hides that.

The finding to look for — and to publish either way:

If volume and volatility gradients agree with each other while both conflict with the return gradient, the correct conclusion is that the two auxiliary tasks form a coherent block that the return task does not belong to — and the fix is task grouping, not more capacity. When conflict is real, the standard remedies are PCGrad (Yu et al., 2020), which projects each conflicting gradient onto the normal plane of the other; CAGrad (Liu et al., 2021), which searches for a descent direction that does not hurt any task; or dropping the auxiliary task entirely.

Note what is deliberately absent: a t-SNE plot of the shared representation coloured by target value. It is decorative — the cosine numbers above say everything the embedding would gesture at, and they say it as numbers.

Validation Protocol

The measurement above is worthless under a sloppy protocol, and MTL makes the usual traps worse because there are three targets to leak instead of one.

Real data, not a simulator. The targets must come from actual OHLCV/trade data. A hardcoded GARCH toy generates volatility that is correlated with returns by construction, which is precisely the thing under test — the experiment would be measuring its own generator. If you want a fitted volatility process, GARCH volatility forecasting for crypto fits GARCH(1,1) by maximum likelihood on real BTC/ETH and validates the standardized residuals, and asymmetric GARCH and the leverage effect covers why a Gaussian symmetric-response simulator misstates crypto volatility in the first place. Synthetic data is defensible only when it provides controlled ground truth — a known, author-set task correlation you are trying to recover — which is a different experiment from the one here.

Scalers fit on train only. Fit the feature scaler and all three target scalers inside each training fold and apply to validation; a global fit_transform before splitting leaks test-set moments into training. This exact failure is catalogued in the look-ahead bias taxonomy.

Purged, embargoed walk-forward folds. One 80/20 chronological split cannot distinguish an MTL improvement from a fold effect — that is the entire argument of walk-forward optimization, which shows three splits producing three conclusions. Reuse the expanding-window purged_walk_forward generator from spread modeling with machine learning: it drops a gap of horizon rows on both sides of each boundary, which matters here because overlapping realized-volatility windows leak across the boundary even when the return target does not.

A classical baseline. An MTL net that beats three single-task nets has proved nothing if a per-target gradient-boosting or ridge model beats all four. Fit one model per target with LightGBM or ridge on the same folds and the same features, and report it in the same table.

What Would Make MTL Worth It Here

Conditions under which MTL should win, stated as hypotheses to check against the folds above rather than as a checklist:

• Auxiliary labels are cleaner than the primary label. Volume is directly observed; "expected return" is not. If the return head is mostly fitting noise, gradient signal from the auxiliary heads is the only well-posed part of the objective.
• Training data is limited relative to encoder capacity, so the auxiliary constraint does real regularization work rather than just competing for parameters.
• Inference latency matters and one forward pass beats three.

And the case against, equally testable: if the measured cos_sim(return, ·) values are persistently negative, the shared encoder is being pulled away from the primary task and the auxiliary heads are a tax, not a regularizer.

Conclusion

Returns, volume, and volatility come from the same microstructure, so a shared representation is a reasonable prior — but a prior is not a result. The two things this setup can actually establish are which loss-balancing scheme the data prefers (with the learned weights reported, not just the winner named) and whether the task gradients on the shared encoder agree, measured over training rather than assumed from the fact that the targets are correlated.

If the purged walk-forward folds show the MTL net failing to beat a per-target gradient-boosting model, that is the finding and it gets published as such — the template is the honest negative. A negative result on negative transfer is still a result about negative transfer.