A DeePM Neural Network for Cross-Asset Portfolio Risk Weights
Code Machine Learning for Trading
Summary
This document describes the architecture of a neural policy network that maps asset-level time series and context into bounded portfolio risk weights. Its per-asset backbone combines context embeddings, feature-wise modulation, variable selection, an LSTM, and causal temporal attention. It then processes relationships across assets with lagged cross-sectional attention and can optionally apply attention restricted by a macroeconomic adjacency graph.
Static context may encode asset identity, group membership, and trading costs, while masks handle unavailable assets. The output head uses a hyperbolic tangent and applies the asset mask, producing weights between negative and positive one for valid assets. The document explains architectural components and safeguards against attention rows with no valid keys, but provides no training procedure, trading results, benchmark comparison, or evidence that the architecture improves portfolio performance. The description alone therefore supports understanding the model design, not judging its empirical effectiveness.
Key ideas
- Asset context can condition both input features and recurrent hidden states.
- Variable selection combines features using context-dependent weights.
- The model uses recurrent and causal attention layers to represent temporal patterns.
- Lagged cross-sectional attention and an optional graph mask model inter-asset relationships.
- A bounded output head produces signed risk weights while respecting asset availability.
Tags
Full text
# model.py
```py
"""PyTorch implementation of the DeePM deep portfolio manager.
Architecture components:
1. Per-asset temporal backbone (shared weights): FiLM conditioning, variable
selection, LSTM, temporal self-attention.
2. Cross-sectional attention with Directed Delay for causality.
3. Macroeconomic graph prior as adjacency-masked attention.
4. Output: bounded risk weight p_{i,t} in (-1, 1) via tanh.
"""
from __future__ import annotations
import torch
from torch import nn
from .configs import ModelConfig
from .utils import causal_attention_mask
class StaticContextEncoder(nn.Module):
"""Encode per-asset static context (asset id, group id, costs)."""
def __init__(
self,
*,
n_assets: int,
n_groups: int | None,
cfg: ModelConfig,
) -> None:
super().__init__()
self.cfg = cfg
self.asset_emb = nn.Embedding(n_assets, cfg.asset_embedding_dim)
self.group_emb: nn.Embedding | None = None
if cfg.use_group_embedding:
if n_groups is None:
raise ValueError("n_groups required when use_group_embedding is True")
self.group_emb = nn.Embedding(n_groups, cfg.group_embedding_dim)
self.include_cost = cfg.use_cost_in_context
@property
def context_dim(self) -> int:
dim = self.cfg.asset_embedding_dim
if self.group_emb is not None:
dim += self.cfg.group_embedding_dim
if self.include_cost:
dim += 1
return dim
def forward(
self,
*,
asset_ids: torch.Tensor,
group_ids: torch.Tensor | None,
costs: torch.Tensor | None,
) -> torch.Tensor:
"""Return context embedding of shape (N, C)."""
emb_list = [self.asset_emb(asset_ids)]
if self.group_emb is not None:
if group_ids is None:
raise ValueError("group_ids required when group_emb is enabled")
emb_list.append(self.group_emb(group_ids))
if self.include_cost:
if costs is None:
raise ValueError("costs required when use_cost_in_context is True")
emb_list.append(costs)
return torch.cat(emb_list, dim=-1)
class FiLM(nn.Module):
"""Feature-wise linear modulation: x -> x * (1 + gamma) + beta."""
def __init__(self, *, context_dim: int, n_features: int) -> None:
super().__init__()
self.proj = nn.Linear(context_dim, 2 * n_features)
self.n_features = n_features
def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
gb = self.proj(context) # (N, 2F)
gamma, beta = gb[:, : self.n_features], gb[:, self.n_features :]
gamma = gamma.unsqueeze(0).unsqueeze(0) # (1, 1, N, F)
beta = beta.unsqueeze(0).unsqueeze(0)
return x * (1.0 + gamma) + beta
class VectorizedVariableSelection(nn.Module):
"""Lightweight variable selection network (V-VSN)."""
def __init__(
self,
*,
n_features: int,
d_model: int,
context_dim: int,
hidden_dim: int,
dropout: float,
) -> None:
super().__init__()
self.n_features = n_features
self.d_model = d_model
self.feature_weight = nn.Parameter(torch.empty(n_features, d_model))
self.feature_bias = nn.Parameter(torch.zeros(n_features, d_model))
nn.init.xavier_uniform_(self.feature_weight)
self.selector = nn.Sequential(
nn.Linear(n_features + context_dim, hidden_dim),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, n_features),
)
self.out_norm = nn.LayerNorm(d_model)
def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
b, t, n, f = x.shape
context_bt = context.unsqueeze(0).unsqueeze(0).expand(b, t, n, -1)
logits = self.selector(torch.cat([x, context_bt], dim=-1))
weights = torch.softmax(logits, dim=-1)
z = torch.einsum("btnf,fd->btnfd", x, self.feature_weight) + self.feature_bias
h = (weights.unsqueeze(-1) * z).sum(dim=-2)
return self.out_norm(h)
class AdapterBlock(nn.Module):
"""FFN adapter with residual connection and LayerNorm."""
def __init__(self, *, d_model: int, hidden_mult: int, dropout: float) -> None:
super().__init__()
d_ff = int(hidden_mult * d_model)
self.ln = nn.LayerNorm(d_model)
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model)
)
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return x + self.dropout(self.ff(self.ln(x)))
class TemporalSelfAttentionBlock(nn.Module):
"""Causal temporal self-attention per asset."""
def __init__(self, *, d_model: int, n_heads: int, dropout: float, adapter_mult: int) -> None:
super().__init__()
self.mha = nn.MultiheadAttention(
embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
)
self.ln = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
self.adapter = AdapterBlock(d_model=d_model, hidden_mult=adapter_mult, dropout=dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
t = x.shape[1]
attn_mask = causal_attention_mask(t, device=x.device)
y, _ = self.mha(x, x, x, attn_mask=attn_mask)
x = self.ln(x + self.dropout(y))
return self.adapter(x)
class CrossSectionalAttention(nn.Module):
"""Cross-asset attention with Directed Delay (time lag)."""
def __init__(self, *, d_model: int, n_heads: int, dropout: float, lag: int) -> None:
super().__init__()
self.lag = int(lag)
self.mha = nn.MultiheadAttention(
embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
)
self.ln = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
b, t, n, d = h.shape
if self.lag > 0:
pad = torch.zeros((b, self.lag, n, d), device=h.device, dtype=h.dtype)
h_kv = torch.cat([pad, h[:, : t - self.lag, :, :]], dim=1)
pad_m = torch.zeros((b, self.lag, n), device=mask.device, dtype=mask.dtype)
m_kv = torch.cat([pad_m, mask[:, : t - self.lag, :]], dim=1)
else:
h_kv = h
m_kv = mask
q = h.reshape(b * t, n, d)
kv = h_kv.reshape(b * t, n, d)
key_padding_mask = m_kv.reshape(b * t, n) < 0.5
# When lag > 0, the first `lag` timesteps have all-zero keys and masks.
# All keys masked → softmax(all -inf) → NaN. Unmask all positions for
# those rows; attention over zero-valued keys yields zero, so the
# residual connection passes through h unchanged.
all_masked = key_padding_mask.all(dim=-1, keepdim=True)
if all_masked.any():
key_padding_mask = key_padding_mask & ~all_masked
out, _ = self.mha(q, kv, kv, key_padding_mask=key_padding_mask)
out = out.reshape(b, t, n, d)
return self.ln(h + self.dropout(out))
class MacroGraphAttention(nn.Module):
"""Adjacency-masked cross-asset attention (GAT-like)."""
def __init__(
self,
*,
d_model: int,
n_heads: int,
dropout: float,
adjacency_mask: torch.Tensor,
) -> None:
super().__init__()
self.register_buffer("adjacency_mask", adjacency_mask.to(dtype=torch.bool))
self.mha = nn.MultiheadAttention(
embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
)
self.ln = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
b, t, n, d = h.shape
x = h.reshape(b * t, n, d)
key_padding_mask = mask.reshape(b * t, n) < 0.5
# Guard: if adjacency_mask + key_padding_mask blocks ALL keys for any
# query, softmax produces NaN. Unmask everything for those rows.
combined = self.adjacency_mask.unsqueeze(0) | key_padding_mask.unsqueeze(1)
all_blocked = combined.all(dim=-1) # (B*T, N) — True if query i has no valid key
if all_blocked.any():
key_padding_mask = key_padding_mask & ~all_blocked
out, _ = self.mha(x, x, x, attn_mask=self.adjacency_mask, key_padding_mask=key_padding_mask)
out = out.reshape(b, t, n, d)
return self.ln(h + self.dropout(out))
class DeepmPolicy(nn.Module):
"""DeePM policy network that outputs risk weights p_{i,t} in (-1, 1)."""
def __init__(
self,
*,
n_assets: int,
n_features: int,
n_groups: int | None,
adjacency_mask: torch.Tensor | None,
cfg: ModelConfig,
) -> None:
super().__init__()
self.n_assets = int(n_assets)
self.n_features = int(n_features)
self.cfg = cfg
self.context_encoder = StaticContextEncoder(n_assets=n_assets, n_groups=n_groups, cfg=cfg)
context_dim = self.context_encoder.context_dim
self.film = FiLM(context_dim=context_dim, n_features=n_features)
self.vvsn = VectorizedVariableSelection(
n_features=n_features,
d_model=cfg.d_model,
context_dim=context_dim,
hidden_dim=cfg.vvsn_hidden_dim,
dropout=cfg.dropout,
)
self.lstm = nn.LSTM(
input_size=cfg.d_model,
hidden_size=cfg.d_model,
num_layers=cfg.lstm_layers,
batch_first=True,
dropout=cfg.dropout if cfg.lstm_layers > 1 else 0.0,
)
self.h0_proj = nn.Linear(context_dim, cfg.lstm_layers * cfg.d_model)
self.c0_proj = nn.Linear(context_dim, cfg.lstm_layers * cfg.d_model)
self.temporal_blocks = nn.ModuleList(
[
TemporalSelfAttentionBlock(
d_model=cfg.d_model,
n_heads=cfg.n_heads,
dropout=cfg.dropout,
adapter_mult=cfg.adapter_hidden_mult,
)
for _ in range(cfg.temporal_mha_layers)
]
)
self.cross_attn = CrossSectionalAttention(
d_model=cfg.d_model,
n_heads=cfg.cross_attention_heads,
dropout=cfg.dropout,
lag=cfg.cross_attention_lag,
)
self.macro_graph: MacroGraphAttention | None = None
if adjacency_mask is not None:
self.macro_graph = MacroGraphAttention(
d_model=cfg.d_model,
n_heads=cfg.macro_gnn_heads,
dropout=cfg.dropout,
adjacency_mask=adjacency_mask,
)
self.head = nn.Linear(cfg.d_model, 1)
def forward(
self,
x: torch.Tensor,
*,
mask: torch.Tensor,
asset_ids: torch.Tensor,
group_ids: torch.Tensor | None,
costs: torch.Tensor | None,
) -> torch.Tensor:
"""Forward pass: features (B,T,N,F) -> risk weights (B,T,N) in (-1,1)."""
b, t, n, f = x.shape
context = self.context_encoder(asset_ids=asset_ids, group_ids=group_ids, costs=costs)
x_mod = self.film(x, context)
h = self.vvsn(x_mod, context) # (B,T,N,D)
# Per-asset temporal backbone
h_bn = h.permute(0, 2, 1, 3).reshape(b * n, t, self.cfg.d_model)
ctx_bn = context.unsqueeze(0).expand(b, n, -1).reshape(b * n, -1)
h0 = (
self.h0_proj(ctx_bn)
.reshape(b * n, self.cfg.lstm_layers, self.cfg.d_model)
.permute(1, 0, 2)
.contiguous()
)
c0 = (
self.c0_proj(ctx_bn)
.reshape(b * n, self.cfg.lstm_layers, self.cfg.d_model)
.permute(1, 0, 2)
.contiguous()
)
h_bn, _ = self.lstm(h_bn, (h0, c0))
for block in self.temporal_blocks:
h_bn = block(h_bn)
h = h_bn.reshape(b, n, t, self.cfg.d_model).permute(0, 2, 1, 3).contiguous()
# Cross-sectional blocks
h = self.cross_attn(h, mask)
if self.macro_graph is not None:
h = self.macro_graph(h, mask)
# Output head -> tanh
p = torch.tanh(self.head(h).squeeze(-1))
return p * mask
```Shown in full with attribution under the source's licence. Licence: MIT
This summary was written by Stratmill's research agent from the original; it is not a copy of the source.