Models¶

CrystaLLM-pi supports two conditioning mechanisms for crystal structure generation: prefix conditioning and residual conditioning. The tokenizer and datasets used to train these models are documented on the Data page.

Each conditional model family includes a configuration, an encoder that maps the conditioning inputs to the representation expected by the transformer, and a GPT-2 model that applies the conditioning.

Tensor shapes and forward() behaviour are documented in the source docstrings and are available through the linked source or help().

The earlier PKV and Slider classes used in the paper are retained for loading released checkpoints. They implement earlier versions of the same conditioning mechanisms and are not part of the current model API, so they are not included below.

Prefix¶

Prefix conditioning represents scalar properties as prefix key-values for GPT-2.

PrefixGPT2Config ¶

Bases: GPT2Config

Configuration for prefix-conditioned GPT-2.

n_input_vector specifies the number of scalar properties, n_prefix_tokens the number of virtual tokens emitted per layer, and n_hidden_cond the conditioning encoder width. Cross-attention is disabled because conditioning enters through cached key-values.

Source code in _models/Prefix_model.py

class PrefixGPT2Config(GPT2Config):
    """Configuration for prefix-conditioned GPT-2.

    `n_input_vector` specifies the number of scalar properties, `n_prefix_tokens` the number of
    virtual tokens emitted per layer, and `n_hidden_cond` the conditioning encoder width.
    Cross-attention is disabled because conditioning enters through cached key-values.
    """

    def __init__(
        self,
        n_input_vector: int = 2,
        n_prefix_tokens: int = 2,
        n_hidden_cond: int = 128,
        dropout: float = 0.1,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.add_cross_attention = False

        self.n_input_vector = n_input_vector
        self.n_prefix_tokens = n_prefix_tokens
        self.n_hidden_cond = n_hidden_cond
        self.dropout = dropout

PrefixEncoder ¶

Bases: Module

Project a scalar conditioning vector into per-layer prefix key-value tensors.

The MLP emits one flat tensor per batch item, which is later reshaped into layers, heads, and prefix tokens. This encoder uses GELU, whereas the legacy PKV_model uses ReLU, so the two checkpoint families are not interchangeable.

Source code in _models/Prefix_model.py

class PrefixEncoder(nn.Module):
    """Project a scalar conditioning vector into per-layer prefix key-value tensors.

    The MLP emits one flat tensor per batch item, which is later reshaped into layers, heads, and
    prefix tokens. This encoder uses GELU, whereas the legacy `PKV_model` uses ReLU, so the two
    checkpoint families are not interchangeable.
    """

    def __init__(self, config: PrefixGPT2Config) -> None:
        super().__init__()
        self.config = config

        # Calculate KV size for every layer.
        self.kv_size = (
            config.n_prefix_tokens
            * config.n_head
            * (config.hidden_size // config.n_head)
            * 2
            * config.n_layer
        )

        self.processor = nn.Sequential(
            nn.Linear(config.n_input_vector, config.n_hidden_cond),
            nn.LayerNorm(config.n_hidden_cond),
            # nn.ReLU(),
            # now GELU
            nn.GELU(),
            nn.Linear(config.n_hidden_cond, config.n_hidden_cond * 2)
        )

        self.to_kv = nn.Linear(config.n_hidden_cond * 2, self.kv_size)
        self.dropout = nn.Dropout(config.dropout) # Use conditioning dropout

    def forward(self, x: torch.Tensor) -> tuple:
        """Encode scalar conditioning values into GPT-2 ``past_key_values``.

        Args:
            x: Scalar conditioning values with shape ``[B, n_input_vector]``.

        Returns:
            One ``(key, value)`` tuple per layer, each with shape
            ``[B, n_head, n_prefix_tokens, head_dim]``.
        """
        batch_size = x.shape[0]

        # The encoder emits one flat prefix blob per batch item. Reshape happens
        # only after the final linear layer so the MLP can stay token-agnostic.
        x = self.processor(x)
        x = self.to_kv(x)
        x = self.dropout(x)
        return reshape_prefix_kv_to_past_key_values(
            x, batch_size, self.config.n_prefix_tokens, self.config
        )

PrefixGPT ¶

Bases: GPT2LMHeadModel

GPT-2 conditioned on scalar properties through learned prefix key-values.

The encoder output occupies config.n_prefix_tokens cached positions before the text sequence. Generation must therefore account for the prefix when computing the effective text length and position embeddings.

Source code in _models/Prefix_model.py

class PrefixGPT(GPT2LMHeadModel):
    """GPT-2 conditioned on scalar properties through learned prefix key-values.

    The encoder output occupies `config.n_prefix_tokens` cached positions before the text sequence.
    Generation must therefore account for the prefix when computing the effective text length and
    position embeddings.
    """
    config_class = PrefixGPT2Config

    def __init__(self, config: PrefixGPT2Config) -> None:
        super().__init__(config)
        self.conditioning = PrefixEncoder(config)

    def forward(
        self,
        input_ids: torch.Tensor | None=None,
        attention_mask: torch.Tensor | None=None,
        condition_values: torch.Tensor | None=None,
        labels: torch.Tensor | None=None,
        **kwargs: object
    ) -> "CausalLMOutputWithCrossAttentions":
        # Check if cached past_key_values exist in kwargs
        """Run GPT-2 with scalar conditioning represented as prefix key-values.

        A ``ValueError`` is raised when neither ``condition_values`` nor cached
        ``past_key_values`` is supplied, because no prefix can then be constructed.

        Args:
            input_ids: Token IDs with shape ``[B, T]``.
            attention_mask: Text attention mask with shape ``[B, T]``. Prefix positions
                are added internally.
            condition_values: Scalar conditioning with shape ``[B, n_input_vector]``.
                Values may include ``MISSING_CONDITION_VALUE``.
            labels: Optional targets with shape ``[B, T]``.

        Returns:
            ``CausalLMOutputWithCrossAttentions`` containing logits with shape
            ``[B, T, vocab_size]``.
        """
        past_key_values = kwargs.pop("past_key_values", None)
        # generate() hands an empty Cache on the first step: treat it as "no prefix yet"
        if past_key_values is None or (isinstance(past_key_values, Cache) and past_key_values.get_seq_length() == 0):
            if condition_values is not None:
                # Prefix tokens are prepended to the attention mask so cached
                # positions remain visible to every decoder layer.
                past_key_values = DynamicCache(self.conditioning(condition_values))
                attention_mask = prepend_prefix_attention_mask(
                    attention_mask, condition_values.shape[0], self.config.n_prefix_tokens
                )
            else:
                past_key_values = None

        if past_key_values is None:
            raise ValueError(
                "WARNING: Prefix Condition values activated but not passed correctly."
            )

        # generate() keeps a text-only mask on cached steps, so cover the prefix slots too
        if (attention_mask is not None and input_ids is not None and isinstance(past_key_values, Cache)
                and attention_mask.shape[1] < past_key_values.get_seq_length() + input_ids.shape[1]):
            attention_mask = torch.cat(
                [attention_mask.new_ones(attention_mask.shape[0], self.config.n_prefix_tokens), attention_mask], dim=1
            )

        # generate() counts cache_position over text only, drop it so GPT-2 offsets positions past the prefix
        kwargs.pop("cache_position", None)

        output = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            labels=labels,
            **kwargs
        )

        return output

Residual¶

Residual conditioning adds the conditioning signal to the attention blocks as a residual contribution.

ResidualGPT2Config ¶

Bases: GPT2Config

Configuration for GPT-2 with residual scalar conditioning.

slider_on enables the conditioning path. slider_n_variables specifies the number of scalar properties, slider_n_hidden the encoder MLP width, and slider_n_heads_sharing_slider the number of attention heads sharing each conditioning projection. The sharing count must divide the number of attention heads.