quantlab.model.predefined.master.AttentionBlock

class quantlab.model.predefined.master.AttentionBlock(d_model, nhead, dropout, across_symbols)

Bases: Module

MASTER’s pre-LayerNorm multi-head attention block with a feed-forward layer.

x is [S_t, N, D]. With across_symbols=False (MASTER’s TAttention) each symbol attends over its own N steps and the scores are not scaled; with across_symbols=True (SAttention) the symbols attend to each other at every step and the scores are divided by sqrt(D / nhead). Both return norm2(LN(x) + MHA(LN(x))) plus the feed-forward layer of that sum.

Examples

>>> block = AttentionBlock(d_model=8, nhead=2, dropout=0.0, across_symbols=True)
>>> block(torch.randn(5, 3, 8)).shape
torch.Size([5, 3, 8])

forward(x)

Attend over steps or over symbols; [S_t, N, D] in and out.

Examples

>>> block.forward(torch.randn(2, 4, 8)).shape
torch.Size([2, 4, 8])