puzzle_emb_len

[bumi001]

I am wondering why puzzle_emb_len is set to 16. HRM uses a value of 1 as far as I can figure out. I see that positions 1 through puzzle_emb_len-1 don't participate in lm_head or q_head computations. However, they participate in the z_L and z_H computations through the attention model. I am trying to understand the reason for this design.

[olivkoch]

All puzzle embeddings participate in lm_head and q_head, but indirectly. During the input stage, puzzle embeddings are concatenated with the input embeddings. This is then fed into the reasoning module, which output (z_H) is then fed into lm_head and q_head. After lm_head processing, only the output are kept (hence the self.lm_head(z_H)[:, self.puzzle_emb_len :] slicing, which removes the processed puzzle embeddings).

The q_head leverages the puzzle embeddings by using the first token embedding as a "CLS token". This tokens attends to all input through attention, therefore the q_head has access to the full puzzle embeddings + input to make the halting decision.

All this means that the full puzzle embeddings influence the output of lm_head and q_head.

[bumi001]

@olivkoch, Thank you for your reply. I understand and agree with your explanation. My question is more about the choice between puzzle_emb_len =1 vs puzzle_emb_len = 16. If you look at the _input_embeddings function, you will see the following code:

            if pad_count > 0:
                puzzle_embedding = F.pad(puzzle_embedding, (0, pad_count))

For puzzle_emb_len = 16, the above code pads 0's for sequence positions 1 through 15 as the input embedding. This is done for every puzzle ids. This will make the q, k, and v vectors for positions 1 through 15 at the very first layer receive contribution only through position embedding.

My question is what is the rationale behind this choice of puzzle_emb_len = 16? Why not use 256 or some other number?

[olivkoch]

I get your question now. Here is my interpretation. puzzle_emb_len must remain relatively low because it is a multiplier in the embeddings manipulation to create "scratchpad tokens":

In _input_embeddings(): pad_count = self.puzzle_emb_len * self.config.hidden_size - puzzle_embedding.shape[-1]

which creates 1 ID Token + (N-1) Compute Tokens.

If puzzle_emb_len is 16, this creates 15 additional embeddings of size hidden_size for the attention layers to work with. That's plenty of embedding space. Using 256 or some other numbers would be way to big.

[tjjd4]

I am curious about this question too, is "scratchpad tokens" talking about concept in LLM chain of thought?
cuz I think TRM model going through recursive maybe is more about refinement but not reasoning (base on inference-time scalability), how will it act as scratchpad tokens? I mean how will it help the calculation process?