Random weights during forward

[XinDongol]

In thoughts_forward (

latent-gemma/src/latent_gemma/main.py

Line 54 in b0260a2

def thoughts_forward(self, num_thoughts, thought_ids, thought_mask, num_of_thought_tokens = 1):

), there are some 'random layers'.

For instance, for each forward pass, thought_proj has random weights:

thought_proj = nn.Sequential(
    nn.Linear(last_hidden.shape[-1], self.config.hidden_size // 2),
    nn.LayerNorm(self.config.hidden_size // 2),
    nn.GELU()
).to(last_hidden.device)

In addition, embed_proj also has random weights:

embed_proj = nn.Linear(
    self.config.hidden_size // 2,
    self.config.hidden_size,
    device=projected_thought.device
)

Are they designed on purpose?

[vicksEmmanuel]

Sure, here is the improved grammar while keeping the Markdown format:

In thoughts_forward (

latent-gemma/src/latent_gemma/main.py

Line 54 in b0260a2

def thoughts_forward(self, num_thoughts, thought_ids, thought_mask, num_of_thought_tokens = 1):

), there are some 'random layers'.
For instance, for each forward pass, thought_proj has random weights:

thought_proj = nn.Sequential(
    nn.Linear(last_hidden.shape[-1], self.config.hidden_size // 2),
    nn.LayerNorm(self.config.hidden_size // 2),
    nn.GELU()
).to(last_hidden.device)

In addition, embed_proj also has random weights:

embed_proj = nn.Linear(
    self.config.hidden_size // 2,
    self.config.hidden_size,
    device=projected_thought.device
)

Are they designed on purpose?

Yes. This approach introduces controlled randomness into the latent reasoning process, which can help prevent the model from falling into fixed thought patterns and encourage more diverse reasoning paths