quaterion_models.heads.softmax_head module

class SoftmaxEmbeddingsHead(output_groups: int, output_size_per_group: int, input_embedding_size: int, dropout: float = 0.0, **kwargs)

Bases: EncoderHead

Provides a concatenation of the independent softmax embeddings groups as a head layer

Useful for deriving embedding confidence.

Schema:

         ┌──────────────────┐
         │    Encoder       │
         └──┬───────────┬───┘
            │           │
            │           │
┌───────────┼───────────┼───────────┐
│           │           │           │
│ ┌─────────▼──┐     ┌──▼─────────┐ │
│ │  Linear    │ ... │  Linear    │ │
│ └─────┬──────┘     └─────┬──────┘ │
│       │                  │        │
│ ┌─────┴──────┐     ┌─────┴──────┐ │
│ │  SoftMax   │ ... │  SoftMax   │ │
│ └─────┬──────┘     └─────┬──────┘ │
│       │                  │        │
│  ┌────┴──────────────────┴─────┐  │
│  │       Concatenation         │  │
│  └──────────────┬──────────────┘  │
│                 │                 │
└─────────────────┼─────────────────┘
                  │
                  ▼

get_config_dict() → Dict[str, Any]

Constructs savable params dict

Returns:

Serializable parameters for __init__ of the Module

transform(input_vectors: Tensor)

Parameters:

input_vectors – shape: (batch_size, …, input_dim)

Returns:

shape (batch_size, …, self.output_size_per_group * self.output_groups)

property output_size: int

training: bool