micm_nlp.models.xpe.encoder¶

The CrossPromptEncoder module (XPE core).

The method from Cross-Prompt Encoder for Low-Performing Languages (Findings of IJCNLP-AACL 2025, arXiv:2508.10352).

XPE, SPT and DUAL are all this one class, separated only by encoder_ratio – the fraction of virtual tokens routed through the encoder. 0 is plain soft prompt tuning, 1 is pure XPE, anything between is DUAL. Code gated on isinstance(pe, CrossPromptEncoder) therefore fires for soft prompt tuning too.

Classes¶

CrossPromptEncoder

Generates the virtual-token embeddings prepended to a model's input.

Module Contents¶

class micm_nlp.models.xpe.encoder.CrossPromptEncoder(config: micm_nlp.models.xpe.config.CrossPromptEncoderConfig)¶

Bases: torch.nn.Module

Generates the virtual-token embeddings prepended to a model’s input.

One class covers all three methods, separated only by encoder_ratio: at 0 the tokens come straight from embedding (plain soft prompt tuning), at 1 they all pass through xpe_embedding and xpe_head (XPE), and in between the two sets are concatenated (DUAL).

Parameters:: config – a CrossPromptEncoderConfig.

Input shape (batch_size, total_virtual_tokens); output shape (batch_size, total_virtual_tokens, token_dim).

Example

from micm_nlp.models.xpe.config import CrossPromptEncoderConfig
from micm_nlp.models.xpe.encoder import CrossPromptEncoder

config = CrossPromptEncoderConfig(
    peft_type='XPE',
    num_virtual_tokens=20,
    token_dim=768,
    encoder_num_layers=12,
    encoder_reparameterization_type='MLP',
    encoder_hidden_size=768,
    encoder_ratio=1.0,
)
prompt_encoder = CrossPromptEncoder(config)

Attributes

embedding (torch.nn.Embedding): The soft-prompt table, used for the tokens that skip the encoder.
xpe_embedding (torch.nn.Embedding): The table feeding the encoder head.
xpe_head (torch.nn.Module): The reparameterization head – MLP, LSTM or attention, per encoder_reparameterization_type.
token_dim (int): Hidden width of the base transformer, and so of the produced embeddings.
input_size (int): Width of the encoder’s input, defaulting to token_dim.
output_size (int): Width of the encoder’s output; always token_dim.
hidden_size (int): Hidden width inside the head.
total_virtual_tokens (int): Virtual tokens across all transformer submodules.
encoder_type: The reparameterization type, as CrossPromptEncoderReparameterizationType.

Build the embeddings and, unless this is plain soft prompt tuning, the head.

Both are always constructed, even when pretrained weights will be loaded over them, so the module is never left in a half-initialised state by a checkpoint that happens to be missing a tensor.

Parameters:: config – the encoder config; encoder_ratio decides how the virtual tokens are split between the plain embedding table and the encoded one.

forward(indices, task_ids=None)¶

Forward pass for the PromptEncoder.

Parameters:

indices (torch.Tensor) – Indices of virtual tokens. Shape: (batch_size, num_virtual_tokens)
task_ids (torch.Tensor) – Task IDs for each example in the batch. Shape: (batch_size,)

get_device()¶

Device the encoder’s parameters live on.

Reads whichever embedding table exists – at encoder_ratio 0 or 1 only one of the two is created.

init_embeddings(num: int, dim: int)¶

Create an embedding table, initialised as encoder_embedding_init_type asks.

Parameters:

num – number of virtual tokens.
dim – embedding width.

Raises:

ValueError – on an unknown initialisation type – silently falling back to the default would be indistinguishable from it working.

load_pretrained_state()¶: Load pretrained weights for embeddings and encoder heads (MLP or LSTM). Initialize embedding parameters only.

normalize_embeddings()¶

Renormalise the embedding rows in place, and report their mean norm.

Applies whatever encoder_embedding_normalize asks for, under no_grad. Note the row filter is by name: any 2-D trainable parameter whose name contains embedding is normalised, which includes the plain soft-prompt table – so this is not an XPE-only operation.

Returns:: the mean row norm after normalising, for logging.

print_all_layers()¶: Print all layers in the model, grouped by trainable and non-trainable.

print_trainable_layers()¶: Print all trainable layer names in the model.

set_grad_requirements()¶

Freeze or unfreeze the embeddings and the head, per config.

original_module parameters are skipped: those are PEFT’s copies of the base model’s weights, which must stay frozen either way. The head is only touched when there is one (encoder_ratio > 0).

embedding_freeze¶

embedding_init_type¶

embedding_normalize¶

embedding_normalize_max_norm¶

encoder_dropout¶

encoder_freeze¶

encoder_num_layers¶

encoder_ratio¶

encoder_type¶

hidden_size¶

init_state_dict_path¶

input_size¶

num_heads¶

num_transformer_submodules¶

output_size¶

spt_virtual_tokens = 0¶

token_dim¶

total_virtual_tokens¶

xpe_virtual_tokens = 0¶