Confusion about the weight gi

[TreasureHunter]

Hi, excited to see your new work, but I have some questions.

From your code

MoH/MoH-LLaMA3/Inference_code/modeling_llama.py

Lines 777 to 836 in eb01368

	query_states = self.q_proj(hidden_states)
	key_states = self.k_proj(hidden_states)
	value_states = self.v_proj(hidden_states)
	gate = query_states.view(-1, self.num_heads, self.head_dim)
	logits = torch.stack([torch.norm(gate[:,i,:], p=2, dim=1) for i in range(self.shared_head, self.num_heads)], dim=1)
	logits_std = logits.std(dim=1, keepdim=True)
	logits_norm = logits / (logits_std / 1)
	gates = F.softmax(logits_norm, dim=1)
	num_tokens, num_experts = gates.shape
	_, indices = torch.topk(gates, k=self.routed_head, dim=1)
	mask = F.one_hot(indices, num_classes=num_experts).sum(dim=1)
	mask = mask.reshape(bsz, q_len, -1)
	mask = torch.cat([torch.ones_like(mask).to(query_states.dtype), mask], dim=-1).to(query_states.dtype)
	query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
	key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
	value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
	kv_seq_len = key_states.shape[-2]
	if past_key_value is not None:
	kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
	cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
	query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
	if past_key_value is not None:
	cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
	key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
	key_states = repeat_kv(key_states, self.num_key_value_groups)
	value_states = repeat_kv(value_states, self.num_key_value_groups)
	if attention_mask is not None:
	if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
	raise ValueError(
	f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
	)
	# SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
	# Reference: https://github.com/pytorch/pytorch/issues/112577.
	if query_states.device.type == "cuda" and attention_mask is not None:
	query_states = query_states.contiguous()
	key_states = key_states.contiguous()
	value_states = value_states.contiguous()
	attn_output = torch.nn.functional.scaled_dot_product_attention(
	query_states,
	key_states,
	value_states,
	attn_mask=attention_mask,
	dropout_p=self.attention_dropout if self.training else 0.0,
	# The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
	is_causal=self.is_causal and attention_mask is None and q_len > 1,
	)
	attn_output = attn_output.transpose(1, 2).contiguous()
	attn_output = attn_output.reshape(bsz, q_len, self.num_heads, self.head_dim)
	attn_output = torch.einsum("bne,bned->bned", mask, attn_output)

In Section 3.2 of your paper, you mention that each attention head is associated with a weight gi calculated through Equation 5. However, in the provided code implementation, it appears that 50% of the heads are designated as shared heads and 25% as routed heads, with the corresponding gi values for these 75% of heads seemingly fixed at 1.

This implementation detail appears inconsistent with the theoretical framework described by the formula. Specifically, I could not locate code segments that align with the mathematical derivation presented in the paper. Could you clarify this discrepancy between the theoretical formulation and its practical implementation?

Looking forward to your reply!

[jpthu17]

LLaMA3-8B is a pre-trained model trained on 15T tokens, while we had only 400B tokens (about 3% of the pre-training data) for continue-tuning. This limited our ability to alter the model structure significantly. To address this, we adjusted our original formula to maximize weight reuse from the pre-trained model while preserving its output distribution.

We observe that weighting the attention head outputs significantly alters the distribution of the output of the attention layer, which necessitates a large amount of training data to restore the original performance. To tackle this, we quantize the routing score and use the straight-through estimator to back-propagate the gradients through the sparsity function.