Conditioning via AdaLN

Author: kazeevn

scripts/data.py

@@ -285,6 +285,8 @@ def compute_symmetry_sites(
         symmetry_dataset["composition"] = symmetry_dataset.apply(get_composition_from_symmetry_sites, axis=1)
         if "formation_energy_per_atom" in dataset.columns:
             symmetry_dataset['formation_energy_per_atom'] = dataset['formation_energy_per_atom']
+        if "energy_above_hull" in dataset.columns:
+            symmetry_dataset['energy_above_hull'] = dataset['energy_above_hull']
         if "band_gap" in dataset.columns:
             symmetry_dataset['band_gap'] = dataset['band_gap']
         if "log_klat" in dataset.columns:

src/wyckoff_transformer/cascade/dataset.py

@@ -170,7 +170,8 @@ def __init__(
         start_dtype: torch.dtype = torch.int64,
         device: str = "cpu",
         augmented_storage_device: Optional[str] = None,
-        target_name = None):
+        target_name = None,
+        extra_fields: Optional[List[str]] = None):
         """
         Args:
             data (dict[str, Tensor | List[Tensor]]):
@@ -234,6 +235,9 @@ def __init__(
             target_name (Optional[str], optional):
             The key in `data` for the target variable, if any. The corresponding
             value `data[target_name]` should be a `torch.Tensor`. Defaults to None.
+            extra_fields (Optional[List[str]], optional):
+            Optional list of additional keys in `data` to copy into `self.data`
+            (e.g. condition features used for AdaLN). Defaults to None.
         """
         if augmented_fields is None:
             augmented_fields = []
@@ -250,7 +254,11 @@ def __init__(
         self.cascade_index_from_field = {name: i for i, name in enumerate(cascade_order)}
         self.augmented_fields = augmented_fields
         self.data = {name: data[name].type(dtype).to(device) for name in cascade_order if name not in augmented_fields}
-        self.max_sequence_length = next(iter(self.data.values())).size(1)
+        if extra_fields:
+            for k in extra_fields:
+                v = data[k]
+                self.data[k] = v.to(device) if hasattr(v, "to") else v
+        self.max_sequence_length = data[cascade_order[0]].size(1)
         self.masks = {name: torch.tensor(masks[name], dtype=dtype, device=device) for name in cascade_order}
         self.pads = {name: torch.tensor(pads[name], dtype=dtype, device=device) for name in cascade_order}
         self.stops = {name: torch.tensor(stops[name], dtype=dtype, device=device) for name in cascade_order}
@@ -346,7 +354,8 @@ def get_augmentation(
     def get_masked_cascade_data(
         self,
         known_seq_len: int,
-        known_cascade_len: int):
+        known_cascade_len: int,
+        return_chosen_indices: bool = False):
 
         # assert known_seq_len < data.shape[1]
         # assert known_seq_len >= 0
@@ -372,6 +381,9 @@ def get_masked_cascade_data(
                 res.append(torch.cat([
                     cascade_vector[:, :known_seq_len],
                     self.masks[name].expand(cascade_vector.size(0), 1)], dim=1))
+        
+        if return_chosen_indices:
+            return self.start_tokens[target_is_viable], res, target[target_is_viable], target_is_viable
         return self.start_tokens[target_is_viable], res, target[target_is_viable]
 
 

src/wyckoff_transformer/cascade/model.py

@@ -11,6 +11,53 @@
 
 logger = logging.getLogger(__name__)
 
+
+class AdaLNTransformerEncoderLayer(TransformerEncoderLayer):
+    def __init__(self, d_model, nhead, condition_dim=1, layer_norm_eps=1e-5, **kwargs):
+        super().__init__(d_model, nhead, layer_norm_eps=layer_norm_eps, **kwargs)
+        # AdaLN provides the affine; strip it from the parent's LayerNorms.
+        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps, elementwise_affine=False)
+        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps, elementwise_affine=False)
+        # DiT-style (1 + gamma) modulation; zero-init weights/biases => identity at init.
+        self.adaLN_modulation1 = nn.Linear(condition_dim, 2 * d_model)
+        self.adaLN_modulation2 = nn.Linear(condition_dim, 2 * d_model)
+        for m in (self.adaLN_modulation1, self.adaLN_modulation2):
+            nn.init.zeros_(m.weight)
+            nn.init.zeros_(m.bias)
+
+    def forward(self, src, cond, src_mask=None, src_key_padding_mask=None, is_causal=False):
+        gamma1, beta1 = self.adaLN_modulation1(cond).chunk(2, dim=-1)
+        gamma2, beta2 = self.adaLN_modulation2(cond).chunk(2, dim=-1)
+
+        x = src
+        if self.norm_first:
+            x_norm = self.norm1(x) * (1 + gamma1.unsqueeze(1)) + beta1.unsqueeze(1)
+            x = x + self.dropout1(self.self_attn(x_norm, x_norm, x_norm, attn_mask=src_mask,
+                                  key_padding_mask=src_key_padding_mask, need_weights=False, is_causal=is_causal)[0])
+            x_norm2 = self.norm2(x) * (1 + gamma2.unsqueeze(1)) + beta2.unsqueeze(1)
+            x = x + self.dropout2(self.linear2(self.dropout(self.activation(self.linear1(x_norm2)))))
+        else:
+            x2 = self.self_attn(x, x, x, attn_mask=src_mask,
+                                key_padding_mask=src_key_padding_mask, need_weights=False, is_causal=is_causal)[0]
+            x = x + self.dropout1(x2)
+            x = self.norm1(x) * (1 + gamma1.unsqueeze(1)) + beta1.unsqueeze(1)
+            
+            x2 = self.linear2(self.dropout(self.activation(self.linear1(x))))
+            x = x + self.dropout2(x2)
+            x = self.norm2(x) * (1 + gamma2.unsqueeze(1)) + beta2.unsqueeze(1)
+        return x
+
+class AdaLNTransformerEncoder(TransformerEncoder):
+    """`nn.TransformerEncoder` that threads a conditioning tensor to each AdaLN layer."""
+    def forward(self, src, cond, mask=None, src_key_padding_mask=None, is_causal=False):
+        output = src
+        for mod in self.layers:
+            output = mod(output, cond=cond, src_mask=mask,
+                         src_key_padding_mask=src_key_padding_mask, is_causal=is_causal)
+        if self.norm is not None:
+            output = self.norm(output)
+        return output
+
 class SpecialEmbedding(torch.nn.Module):
     ScalarPassThrough = 0
     VectorPassThrough = 1
@@ -199,7 +246,8 @@ def __init__(self,
                  aggregation_weight: Optional[int] = None,
                  emebdding_dropout: Optional[float] = None,
                  prediction_perceptron_dropout: Optional[float] = None,
-                 concat_start_to_prediction_input_embedding_dim: Optional[int] = None):
+                 concat_start_to_prediction_input_embedding_dim: Optional[int] = None,
+                 condition_dim: Optional[int] = None):
         """
         Expects tokens in the following format:
         START_k -> [] -> STOP -> PAD
@@ -250,8 +298,16 @@ def __init__(self,
         if "nhead" in TransformerEncoderLayer_args and self.d_model % TransformerEncoderLayer_args["nhead"]:
             logger.warning("d_model is not divisible by nhead, padding to the next multiple")
             self.d_model += TransformerEncoderLayer_args["nhead"] - self.d_model % TransformerEncoderLayer_args["nhead"]
-        self.encoder_layers = TransformerEncoderLayer(self.d_model, batch_first=True, **TransformerEncoderLayer_args)
-        self.transformer_encoder = TransformerEncoder(self.encoder_layers, **TransformerEncoder_args)
+        
+        self.condition_dim = condition_dim
+        if condition_dim is not None:
+            self.encoder_layers = AdaLNTransformerEncoderLayer(
+                self.d_model, batch_first=True, condition_dim=condition_dim, **TransformerEncoderLayer_args)
+            self.transformer_encoder = AdaLNTransformerEncoder(self.encoder_layers, **TransformerEncoder_args)
+        else:
+            self.encoder_layers = TransformerEncoderLayer(self.d_model, batch_first=True, **TransformerEncoderLayer_args)
+            self.transformer_encoder = TransformerEncoder(self.encoder_layers, **TransformerEncoder_args)
+            
         self.start_type = start_type
         if start_type == "categorial":
             self.start_embedding = nn.Embedding(n_start, self.d_model)
@@ -347,7 +403,8 @@ def forward(self,
                 start: Tensor,
                 cascade: List[Tensor],
                 padding_mask: Tensor|None,
-                prediction_head: int|None) -> Tensor:
+                prediction_head: int|None,
+                cond: Tensor|None = None) -> Tensor:
         """
         Arguments:
             start: Tensor of shape ``[batch_size]`` with the start token.
@@ -356,6 +413,7 @@ def forward(self,
             prediction_head: Index of the prediction head to use. If None, use the only one. The
                 model works in two stages. Firstly, a vector is prepared wih Encoder and
                 various tweaks. Then, the vector is passed to a perceptron aka prediction head.
+            cond: Tensor of shape ``[batch_size, condition_dim]`` with the conditioning vector for AdaLN.
         Returns:
             Tensor of shape ``[batch_size, seq_len, output_dim]`` with the predictions.
         """
@@ -377,7 +435,12 @@ def forward(self,
         data = torch.cat([self.start_embedding(start).unsqueeze(1), cascade_embedding], dim=1)
         logger.debug("Data size: %s", data.size())
         logger.debug("Padding mask size: %s", padding_mask.size() if padding_mask is not None else "None")
-        transformer_output = self.transformer_encoder(data, src_key_padding_mask=padding_mask)
+        if getattr(self, "condition_dim", None) is not None:
+            if cond is None:
+                raise ValueError("condition_dim is set but cond is not provided")
+            transformer_output = self.transformer_encoder(data, src_key_padding_mask=padding_mask, cond=cond)
+        else:
+            transformer_output = self.transformer_encoder(data, src_key_padding_mask=padding_mask)
 
         logging.debug("Transformer output size: %s", transformer_output.size())
         if self.aggregate_after_encoder:

src/wyckoff_transformer/generator.py

@@ -63,13 +63,16 @@ def __init__(self,
         self.token_engineers = token_engineers
 
 
-    def calibrate(self, dataset: AugmentedCascadeDataset, calibration_element_count_threshold: int = 100):
+    def calibrate(self, dataset: AugmentedCascadeDataset, calibration_element_count_threshold: int = 100, condition_feature: Optional[str] = None):
         """
         The calibraiton is going to be per cascade field.
         We will generate p_predicted and p_true for each cascade field for each
         known sequence length.
         """
         assert dataset.cascade_order == self.cascade_order
+
+        full_cond = dataset.data[condition_feature] if condition_feature is not None else None
+
         with torch.no_grad():
             self.model.eval()
             self.calibrators = []
@@ -84,14 +87,17 @@ def calibrate(self, dataset: AugmentedCascadeDataset, calibration_element_count_
                 logging.info("Calibrating cascade field %s", cascade_name)
                 for known_seq_len in range(dataset.max_sequence_length):
                     if known_cascade_len == 0:
-                        start_tokens, masked_data, target = dataset.get_masked_multiclass_cascade_data(
+                        start_tokens, masked_data, target, chosen_indices = dataset.get_masked_multiclass_cascade_data(
                             known_seq_len, known_cascade_len,
-                            target_type=TargetClass.NextToken, multiclass_target=True)
+                            target_type=TargetClass.NextToken, multiclass_target=True,
+                            return_chosen_indices=True)
                     else:
-                        start_tokens, masked_data, target = dataset.get_masked_multiclass_cascade_data(
+                        start_tokens, masked_data, target, chosen_indices = dataset.get_masked_multiclass_cascade_data(
                             known_seq_len, known_cascade_len,
-                            target_type=TargetClass.NextToken, multiclass_target=False)
-                    model_output = self.model(start_tokens, masked_data, None, known_cascade_len)
+                            target_type=TargetClass.NextToken, multiclass_target=False,
+                            return_chosen_indices=True)
+                    iter_cond = full_cond[chosen_indices] if full_cond is not None else None
+                    model_output = self.model(start_tokens, masked_data, None, known_cascade_len, cond=iter_cond)
                     # Enought data for separate calibration
                     if target.size(0) >= calibration_element_count_threshold:
                         # If model_output and target are on different devices, not our problem
@@ -125,6 +131,7 @@ def generate_tensors(
         max_length: Optional[int] = None,
         elements_vocab: Optional[Dict] = None,
         delimiter: str = "-",
+        cond: Optional[Tensor] = None
     ) -> List[Tensor] | Tuple[List[Tensor], List[float], List[float]]:
         """
         Generates a sequence of tokens.
@@ -245,7 +252,7 @@ def _parse_string_to_ids(s: str) -> set:
                 if self.cascade_is_target.get(cascade_name, False):
                     # +1 for MASK
                     this_generation_input = [generated_cascade[:, :known_seq_len + 1] for generated_cascade in generated]
-                    logits = self.model(start, this_generation_input, None, known_cascade_len)
+                    logits = self.model(start, this_generation_input, None, known_cascade_len, cond=cond)
                     if self.calibrators is not None:
                         if known_seq_len < len(self.calibrators[known_cascade_len]):
                             logits = self.calibrators[known_cascade_len][known_seq_len](logits)