micm_nlp.models.xpe.heads

Reparameterization heads used by CrossPromptEncoder.

Classes

LSTMWrapper

Chain an LSTM into an MLP, discarding the hidden state.

LightweightSelfAttentionHead

One pre-norm self-attention block: attention, residual, LayerNorm, projection.

Functions

gen_attn_head(num_heads, input_size, hidden_size, ...)

Build the attention reparameterization head.

gen_lstm_head(input_size, hidden_size, output_size, ...)

Build the LSTM reparameterization head.

gen_mlp_head(input_size, hidden_size, output_size, ...)

Build the MLP reparameterization head.

Module Contents

class micm_nlp.models.xpe.heads.LSTMWrapper(lstm: torch.nn.LSTM, mlp: torch.nn.Module)

Bases: torch.nn.Module

Chain an LSTM into an MLP, discarding the hidden state.

Exists because torch.nn.LSTM returns (output, (h, c)), which torch.nn.Sequential cannot pass on.

Parameters:

  • lstm – the recurrent layer.

  • mlp – projection applied to its output.

forward(x)

Run the LSTM, drop the hidden state, project the output sequence.

lstm
mlp
class micm_nlp.models.xpe.heads.LightweightSelfAttentionHead(num_heads: int, embed_dim: int, output_dim: int, dropout: float = 0.1)

Bases: torch.nn.Module

One pre-norm self-attention block: attention, residual, LayerNorm, projection.

“Lightweight” is literal – a single block with no feed-forward sublayer, which is all the prompt needs. It is the attention variant of the reparameterization head, selected by CrossPromptEncoderReparameterizationType.

Parameters:

  • num_heads – attention heads; must divide embed_dim.

  • embed_dim – width of the incoming virtual-token embeddings.

  • output_dim – width to project to.

  • dropout – dropout on the attention output, before the residual.

forward(x)
Parameters:

x – (batch_size, num_tokens, embed_dim).

Returns:

(batch_size, num_tokens, output_dim).

dropout
layernorm
out_proj
self_attn
micm_nlp.models.xpe.heads.gen_attn_head(num_heads, input_size, hidden_size, output_size, dropout)

Build the attention reparameterization head.

Self-attention over the virtual tokens, then a single-layer MLP. Like the LSTM head this lets each virtual token see the others, but without a recurrence.

Returns:

a torch.nn.Sequential.

micm_nlp.models.xpe.heads.gen_lstm_head(input_size, hidden_size, output_size, num_layers, dropout)

Build the LSTM reparameterization head.

A bidirectional LSTM followed by a single-layer MLP; the MLP’s input is hidden_size * 2 because the two directions are concatenated. Being bidirectional, each virtual token is conditioned on the whole prompt, not only on the tokens before it.

Returns:

an LSTMWrapper – the LSTM cannot be put in a Sequential directly because it returns a tuple.

micm_nlp.models.xpe.heads.gen_mlp_head(input_size, hidden_size, output_size, num_layers, dropout=0.1)

Build the MLP reparameterization head.

num_layers counts Linear layers in total: the input projection, then num_layers - 1 hidden blocks, then the output projection. Each hidden block is Linear/ReLU/Dropout; the output projection has neither activation nor dropout, since it produces embeddings rather than features.

Returns:

a torch.nn.Sequential.