Size Intensive Readouts Using Multihead Attention - Vanishing Gradient

[bismuthadams1]

Hi MACE team, I've been working on a Transformer style readout inspired by the discussion here: https://github.com/ACEsuit/mace/discussions/917. My model is potentially quite ambitious, having both a regressor and classifier forming a multitask learning problem. I've constructed my model constructor for the transformer as follows:

            self.products = torch.nn.ModuleList([prod])

            layer_blocks = torch.nn.ModuleList()
            transformer_blocks_cls = torch.nn.ModuleList()
            transformer_blocks_regress= torch.nn.ModuleList()

            self.readout_cls = readout_cls

            self.readouts_classifier = torch.nn.ModuleList()

            HEAD_CHANNELS = 16

            readout_cls = (
                LinearReadoutBlock(
                        irreps_in=hidden_irreps, 
                        irrep_out=o3.Irreps(f"{HEAD_CHANNELS}x0e")
                    )
            )
            transformer_blocks_cls.append(
                TransformerGraphReadoutBlock(
                        irreps_in=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                        MLP_irreps = o3.Irreps(f"{HEAD_CHANNELS}x0e")
                    )
            )

            self.readouts_regressor = torch.nn.ModuleList()

            readout_regress = (
                LinearReadoutBlock(
                        irreps_in=hidden_irreps, 
                        irrep_out=o3.Irreps(f"{HEAD_CHANNELS}x0e")
                    )
            )
            transformer_blocks_regress.append(
                TransformerGraphReadoutBlock(
                        irreps_in=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                        MLP_irreps = o3.Irreps(f"{HEAD_CHANNELS}x0e")
                )
            )      

            irreps_by_layer = []
            irreps_by_layer.append(hidden_irreps)

            layer_blocks.append(
                TwinReadouts(
                    readout_cls, readout_regress
                )
            )

            for i in range(num_interactions - 1):

                hidden_irreps_out = hidden_irreps
                
                irreps_by_layer.append(hidden_irreps_out)

                inter = interaction_cls(
                    node_attrs_irreps=node_attr_irreps,
                    node_feats_irreps=hidden_irreps,
                    edge_attrs_irreps=sh_irreps,
                    edge_feats_irreps=edge_feats_irreps,
                    target_irreps=interaction_irreps,
                    hidden_irreps=hidden_irreps_out,
                    avg_num_neighbors=avg_num_neighbors,
                    radial_MLP=radial_MLP,
                )
                self.interactions.append(inter)
                prod = EquivariantProductBasisBlock(
                    node_feats_irreps=interaction_irreps,
                    target_irreps=hidden_irreps_out,
                    correlation=correlation,
                    num_elements=num_elements,
                    use_sc=True,
                )
                self.products.append(prod)
                irreps_by_layer.append(hidden_irreps_out)

                if i == num_interactions - 2:
                    # last layer non-linear readout
                    ro_cls = self.readout_cls(
                        hidden_irreps, (HEAD_CHANNELS*MLP_irreps).simplify(),
                        gate=torch.nn.functional.silu,
                        irrep_out=o3.Irreps(f"{HEAD_CHANNELS}x0e"), num_heads=HEAD_CHANNELS,
                    )
                    ro_reg = self.readout_cls(
                        hidden_irreps, (HEAD_CHANNELS*MLP_irreps).simplify(),
                        gate=torch.nn.functional.silu,
                        irrep_out=o3.Irreps(f"{HEAD_CHANNELS}x0e"), num_heads=HEAD_CHANNELS,
                    )
                    transformer_blocks_cls.append(
                        TransformerGraphReadoutBlock(
                            irreps_in=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                            MLP_irreps=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                    ))
                    transformer_blocks_regress.append(
                        TransformerGraphReadoutBlock(
                            irreps_in=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                            MLP_irreps=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                    ))
                else:
                    ro_cls = LinearReadoutBlock(
                        irreps_in=hidden_irreps,
                        irrep_out=o3.Irreps(f"{HEAD_CHANNELS}x0e")
                    )
                    ro_reg = LinearReadoutBlock(
                        irreps_in=hidden_irreps,
                        irrep_out=o3.Irreps(f"{HEAD_CHANNELS}x0e")
                    )
                    transformer_blocks_cls.append(
                        TransformerGraphReadoutBlock(
                            irreps_in=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                            MLP_irreps=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                    ))
                    transformer_blocks_regress.append(
                        TransformerGraphReadoutBlock(
                            irreps_in=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                            MLP_irreps=o3.Irreps(f"{HEAD_CHANNELS}x0e"),
                    ))

                layer_blocks.append(
                     TwinReadouts(ro_cls, ro_reg)
                )

            
            self.readouts = torch.nn.ModuleDict({
                'layers' : layer_blocks,
                'transformers_cls' : transformer_blocks_cls,
                'transformers_regress' : transformer_blocks_regress
            })  # Add to readouts attribute to expose them to optimizer as parameters

And the transformer block itself is:

@compile_mode("script")
class TransformerGraphReadoutBlock(torch.nn.Module):
    def __init__(self, irreps_in: o3.Irreps, MLP_irreps: o3.Irreps, cueq_config=None):
        super().__init__()
        input_dim = irreps_in.num_irreps

        self.mom_mapper = torch.nn.Sequential(
            torch.nn.Linear(3, input_dim),
            torch.nn.GELU(),
            torch.nn.Dropout(0.01),
        )

        mid_dim = MLP_irreps.num_irreps
        self.mom_attn = torch.nn.MultiheadAttention(
            input_dim, 8, 0.05, batch_first=True 
        )

        self.fc = torch.nn.Sequential(
            torch.nn.Linear(input_dim, mid_dim),
            torch.nn.GELU(),
            torch.nn.Dropout(0.01),
            torch.nn.Linear(mid_dim, 1),
        )
    def forward(self, x):   # x [B, C, 3]
        
        B, C, D = x.shape 
        h = self.mom_mapper(x)            # [B, C, d_model] (linear applies on last dim) output B,C,1
        h_attn, _ =  self.mom_attn(h,h,h) # _ contains weights/ 
        h = h  + h_attn                   # Residual
        h = h.mean(dim=1)                         # pool over C -> [B, E]
        h = self.fc(h)

        return h

Following a similar design to Ilyes's suggestion. Unfortunately when training the gradients seem to vanish. It doesn't seem to be an issue with my loss function design or wiring issue as when I replace my the transformer architecture with a more simple scatter_mean over a set of o3.Irreps("1x0e") per layer I see some gradient but the model performance is suboptimal (I suspect because the independent layers of the regressor and classifier aren't very large- my hope was the transformer would give more independent layers to each task).

I hope this makes sense.

[ilyes319]

Hey @bismuthadams1, may i ask what kind of quantities you are trying to fit?
I would recommend you try to following things (in order):

• one common mistake is to forget to add the new layers to the optimizer in MACE. Make sure to actually add it else it will not train:

  mace/mace/tools/scripts_utils.py

  Line 736 in 42cd495

  def get_params_options(

  . MACE uses a custom set of params options, so you need to do that manually!
• add some layer normalization before each block of the readout. You can also try to do x + layernorm(x). You can also try to add a layernorms directly in the self.fc. You can plot the distribution of activations norms in the very few training steps and post it here, and i can tell you how to tune it.(see here for example how to do it in fc

  mace/mace/modules/radial.py

  Line 361 in 42cd495

  class RadialMLP(torch.nn.Module):

  ).
• You can swtich the optimizer, i recommend using the schedulerfree optimizer in MACE as it is the most robust. You might need to play with the beta + weight decay value.
• you can try to change the initialization of the Linear in the mom_mapper and fc to something like Xavier.
• you switch around GELU to SILU.
• the value of the dropout is really task dependent, I recommend you set it to no dropout first and then increase progressively.

[bismuthadams1]

Hey @ilyes319 , thanks for the detailed reply. Regarding what I'm trying to fit, I am using MACE quite unconventionally, and looking at predicting a degree of chemical interaction (I am keeping vague as it is not yet published). The data is quite sparse with lots of null values (no interaction), so I added a classifier to attempt filter them out, and only trained on the positive samples.

I will address some of the points I can address quickly now and get back to you on some of your other suggestions.

• For the model parameters, I ensured everything is added back to the self.readout attribute (I added a class which takes both regress and classifier readouts and add them to a moduledict() which is set to self.readout, so hopefully everything is fixed here. In fact, because I am using a dynamic loss weighting in training (as I'm building a multitask model), I have exposed the loss weights to the optimizer here.
• Thanks for the layer norm suggestion, I will try to implement this. I suspect the x + layernorm(x) will help with the vanishing gradient.
• I will try the scedulerfree option
• I will try change the intialization.
• Currently I am using SILU already, but thanks for this. EDIT: my bad I realized I literally posted GeLU in the code snippet!
• I will remove drop out here and work upwards.