Modulo parameter constraints

1. modulo

hidden_size % num_attention_heads == 0

use case

In transformer models, we use multi-head attention which the hidden vector will be divided into n part, and n is number of attention heads, so we usually want hidden_size divisible by num_attention_heads

2. log2

math.log2(batch_size) % 1 == 0

use case

To make batch size be 2**n to just fit in memory.

3. Why?

Can we just pass a function as parameter constraint and have all parameters’ names as arguments?

[Balandat]

Can you just reparameterize your problem in this case? Like define integer parameters size_per_head and log_batch_size and then in your evaluation code have

hidden_size = size_per_head * num_attention_heads batch_size = 2**log_batch_size

Or would that mean you need a joint constraint on size_per_head and num_attention_heads?

@sdsingh From an optimization/candidate generation perspective (other than difficulty of the problem) there is no issue with imposing non-linear constraints on the parameter space. I assume the linearity assumption is mainly imposing structure for representing the constraints in Ax?

[sdsingh]

All parameter constraints are implemented as linear constraints in the modeling layer. When we pass points into our GPs, they are all normalized to [0,1]^d, where then our constraints are applied. You can see how these are implemented w/ a simple matrix multiply in the Botorch model. As a result, we can only support constraints that can be mapped into a linear constraint, however creatively that may be. Unfortunately, I don’t see a way of doing that for either of these constraint types.

Both of these constraint types seem pretty useful, though. Let me think about what we can do here, but unfortunately I don’t think this will be a quick fix.