Cách Chúng Tôi Huấn Luyện Shannon V1.5 Tư Duy Bằng GRPO

1. Động lực: Tại sao lại là Mô hình tư duy?

Đối vớiđội đỏ AInghiên cứu, việc hiểucáchmột AI đưa ra các đầu ra có khả năng gây hại cũng quan trọng như việc nghiên cứu chính các đầu ra đó. Các mô hình tư duy của Shannon V1.5 phơi bày toàn bộ chuỗi tư duy, cho phép các nhà nghiên cứu:

• Theo dõi đường suy luận dẫn đếnhậu quả AI không kiểm duyệthành vi
• Xác định nơi các rào chắn an toàn thường can thiệp (chứng minhtầm quan trọng của rào chắn an toàn AI)
• Phát triển các cơ chế an toàn mạnh mẽ hơn nhắm mục tiêu vào các mẫu suy luận cụ thể
• Tạo ra các trường hợp thử nghiệm đối kháng tốt hơn dựa trên các chuỗi suy luận đã quan sát

2. Tìm hiểu về GRPO

Tối ưu hóa chính sách tương đối nhóm (GRPO)là một tiến bộ so với các phương pháp RLHF truyền thống, cho phép huấn luyện khả năng suy luận ổn định và hiệu quả hơn. Được phát triển bởi DeepSeek AI, nó đã chứng minh hiệu quả đặc biệt đối với việc huấn luyện chuỗi tư duy.

Tại sao GRPO vượt trội hơn RLHF truyền thống?

Nền tảng toán học của GRPO

GRPO tối ưu hóa chính sách bằng cách so sánh các phản hồi trong các nhóm thay vì so với một mô hình phần thưởng tuyệt đối:

So sánh tương đối này có một số ưu điểm:

• Chuẩn hóa:Tự động điều chỉnh cho độ khó khác nhau giữa các lời nhắc
• Ổn định:Giảm phương sai trong ước tính gradient
• Hiệu quả:Không cần mô hình phần thưởng riêng biệt

def compute_grpo_loss(
    policy_logprobs: torch.Tensor,
    rewards: torch.Tensor,
    group_size: int = 8
) -> torch.Tensor:
    """
    Compute GRPO loss with group-relative reward normalization.
    
    Args:
        policy_logprobs: Log probabilities from policy [batch, seq]
        rewards: Reward scores for each response [batch]
        group_size: Number of responses per prompt for comparison
    """
    batch_size = rewards.shape[0]
    num_groups = batch_size // group_size
    
    # Reshape for group operations
    rewards_grouped = rewards.view(num_groups, group_size)
    logprobs_grouped = policy_logprobs.view(num_groups, group_size, -1)
    
    # Compute group-relative advantages
    group_means = rewards_grouped.mean(dim=1, keepdim=True)
    group_stds = rewards_grouped.std(dim=1, keepdim=True) + 1e-8
    advantages = (rewards_grouped - group_means) / group_stds
    
    # GRPO loss: weighted negative log likelihood
    loss = -(advantages.unsqueeze(-1) * logprobs_grouped).sum(dim=-1).mean()
    
    return loss

3. Chưng cất DeepSeek

Để khởi động khả năng tư duy của Shannon V1.5, chúng tôi đã chưng cất các mẫu chuỗi tư duy từ các mô hình suy luận của DeepSeek. Điều này cung cấp các dấu vết CoT chất lượng cao để huấn luyện đầu tư duy của chúng tôi.

Thành phần tập dữ liệu DeepSeek

Quy trình thu thập dấu vết

Chúng tôi đã thu thập các dấu vết tư duy trên nhiều lĩnh vực đa dạng để đảm bảo phạm vi suy luận toàn diện:

class DeepSeekDistiller:
    """Distill chain-of-thought traces from DeepSeek models."""
    
    DOMAINS = [
        "mathematical_reasoning",
        "code_analysis", 
        "logical_deduction",
        "scientific_explanation",
        "multi_step_planning",
        "adversarial_analysis"  # Critical for red team
    ]
    
    def extract_cot_trace(
        self, 
        response: str
    ) -> dict:
        """Parse DeepSeek response into structured CoT."""
        
        # DeepSeek uses ... tags
        think_match = re.search(
            r'(.*?)', 
            response, 
            re.DOTALL
        )
        
        if not think_match:
            return None
            
        thinking = think_match.group(1)
        final_answer = response.split('')[-1].strip()
        
        # Parse individual reasoning steps
        steps = self.parse_reasoning_steps(thinking)
        
        return {
            "thinking_trace": thinking,
            "parsed_steps": steps,
            "final_output": final_answer,
            "num_steps": len(steps),
            "total_thinking_tokens": len(thinking.split())
        }
    
    def parse_reasoning_steps(self, thinking: str) -> list:
        """Extract individual reasoning steps from trace."""
        # Split on common step indicators
        step_patterns = [
            r'\n\d+\.',           # "1. ", "2. "
            r'\nStep \d+:',       # "Step 1:"
            r'\n(?:First|Next|Then|Finally),',
            r'\n- '              # Bullet points
        ]
        
        combined_pattern = '|'.join(step_patterns)
        steps = re.split(combined_pattern, thinking)
        
        return [s.strip() for s in steps if s.strip()]

4. Kiến trúc đầu tư duy

Các mô hình Shannon V1.5 tích hợp mộtđầu tư duytạo ra các dấu vết suy luận rõ ràng trước khi đưa ra kết quả cuối cùng. Việc bổ sung kiến trúc này cho phép CoT minh bạch mà không cần sửa đổi kiến trúc Mixtral cơ bản.

Triển khai đầu tư duy

class ThinkingHead(nn.Module):
    """
    Dedicated thinking module for Shannon V1.5.
    Generates explicit chain-of-thought traces.
    """
    
    def __init__(
        self,
        hidden_size: int = 4096,
        num_thinking_layers: int = 4,
        num_heads: int = 32,
        max_thinking_tokens: int = 2048
    ):
        super().__init__()
        
        self.hidden_size = hidden_size
        self.max_thinking_tokens = max_thinking_tokens
        
        # Special tokens
        self.think_start = nn.Parameter(torch.randn(1, 1, hidden_size))
        self.think_end = nn.Parameter(torch.randn(1, 1, hidden_size))
        
        # Thinking transformer layers
        self.thinking_layers = nn.ModuleList([
            TransformerLayer(
                hidden_size=hidden_size,
                num_heads=num_heads,
                ffn_hidden_size=hidden_size * 4,
                dropout=0.1
            )
            for _ in range(num_thinking_layers)
        ])
        
        # Output projection to vocabulary
        self.output_proj = nn.Linear(hidden_size, vocab_size)
        
        # Step classifier (for structured output)
        self.step_classifier = nn.Linear(hidden_size, 5)  # 5 step types
    
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        generate_steps: bool = True
    ) -> dict:
        """
        Generate thinking trace from input hidden states.
        
        Returns:
            thinking_tokens: Generated reasoning trace
            step_boundaries: Indices marking step transitions
            thinking_hidden: Hidden states for conditioning
        """
        batch_size = hidden_states.shape[0]
        
        # Prepend thinking start token
        thinking_input = torch.cat([
            self.think_start.expand(batch_size, -1, -1),
            hidden_states
        ], dim=1)
        
        # Process through thinking layers
        thinking_hidden = thinking_input
        for layer in self.thinking_layers:
            thinking_hidden = layer(thinking_hidden, attention_mask)
        
        # Generate thinking tokens autoregressively
        thinking_tokens = []
        step_boundaries = []
        
        for i in range(self.max_thinking_tokens):
            logits = self.output_proj(thinking_hidden[:, -1, :])
            next_token = logits.argmax(dim=-1)
            
            # Check for step boundaries
            step_type = self.step_classifier(thinking_hidden[:, -1, :])
            if step_type.argmax(dim=-1) != 0:  # 0 = continue
                step_boundaries.append(i)
            
            thinking_tokens.append(next_token)
            
            # Check for think_end
            if next_token == self.think_end_token_id:
                break
            
            # Update for next iteration
            # ... (autoregressive generation logic)
        
        return {
            "thinking_tokens": torch.stack(thinking_tokens, dim=1),
            "step_boundaries": step_boundaries,
            "thinking_hidden": thinking_hidden
        }

5. Quy trình huấn luyện

Giai đoạn 1: Huấn luyện trước đầu tư duy

Đầu tiên, chúng tôi huấn luyện trước đầu tư duy trên các dấu vết CoT được chắt lọc từ DeepSeek bằng cách sử dụng hàm mất mát entropy chéo tiêu chuẩn:

# Thinking Head Pre-training Configuration
model:
  base: shannon-ai/v1-deep  # Start from GPT-5 distilled model
  thinking_head:
    num_layers: 4
    hidden_size: 4096
    max_tokens: 2048

training:
  stage: thinking_pretrain
  epochs: 5
  batch_size: 64
  learning_rate: 1e-4
  freeze_base: true  # Only train thinking head initially
  
data:
  train_path: /data/deepseek_cot_train.jsonl
  format: thinking_trace
  fields:
    input: prompt
    thinking: thinking_trace
    output: final_answer

Giai đoạn 2: Tinh chỉnh GRPO

Sau khi huấn luyện trước, chúng tôi áp dụng GRPO để cải thiện chất lượng tư duy bằng cách sử dụng các so sánh tương đối theo nhóm:

class GRPOTrainer:
    """GRPO trainer for thinking model optimization."""
    
    def __init__(
        self,
        model: ThinkingModel,
        group_size: int = 8,
        kl_coef: float = 0.1
    ):
        self.model = model
        self.group_size = group_size
        self.kl_coef = kl_coef
        self.ref_model = copy.deepcopy(model)
        self.ref_model.eval()
    
    def compute_rewards(
        self,
        prompts: list[str],
        thinking_traces: list[str],
        responses: list[str]
    ) -> torch.Tensor:
        """
        Compute rewards for thinking quality.
        Multiple signals combined for comprehensive evaluation.
        """
        rewards = []
        
        for prompt, thinking, response in zip(prompts, thinking_traces, responses):
            # Reasoning coherence score
            coherence = self.evaluate_coherence(thinking)
            
            # Step structure quality
            structure = self.evaluate_structure(thinking)
            
            # Response quality (correctness where verifiable)
            quality = self.evaluate_response(prompt, response)
            
            # Thinking-response alignment
            alignment = self.evaluate_alignment(thinking, response)
            
            # Combined reward
            reward = (
                0.3 * coherence +
                0.2 * structure +
                0.3 * quality +
                0.2 * alignment
            )
            rewards.append(reward)
        
        return torch.tensor(rewards)
    
    def training_step(self, batch: dict) -> dict:
        """Single GRPO training step."""
        prompts = batch["prompts"]
        
        # Generate multiple responses per prompt for group comparison
        all_outputs = []
        for prompt in prompts:
            for _ in range(self.group_size):
                output = self.model.generate_with_thinking(
                    prompt,
                    temperature=0.8,  # Diversity for comparison
                    do_sample=True
                )
                all_outputs.append(output)
        
        # Compute rewards
        rewards = self.compute_rewards(
            prompts=[p for p in prompts for _ in range(self.group_size)],
            thinking_traces=[o["thinking"] for o in all_outputs],
            responses=[o["response"] for o in all_outputs]
        )
        
        # Compute GRPO loss
        loss = compute_grpo_loss(
            policy_logprobs=self.get_logprobs(all_outputs),
            rewards=rewards,
            group_size=self.group_size
        )
        
        # Add KL penalty against reference model
        kl_div = self.compute_kl_divergence(all_outputs)
        total_loss = loss + self.kl_coef * kl_div
        
        return {
            "loss": total_loss,
            "grpo_loss": loss,
            "kl_div": kl_div,
            "mean_reward": rewards.mean()
        }

Giai đoạn 3: Chuyên môn hóa đội đỏ

Cuối cùng, chúng tôi tinh chỉnh thêm trên các kịch bản đối kháng để đảm bảo các dấu vết tư duy phơi bày đúng cách suy luận choAI không kiểm duyệt tiếp theophân tích:

6. Kết quả & Phân tích

Các chỉ số chất lượng tư duy

Những phát hiện chính

• Tính minh bạch được cải thiện đáng kể:Từ 12% lên 97.8% suy luận hiện được diễn đạt rõ ràng
• Độ chính xác của suy luận tăng lên:Tư duy rõ ràng đã cải thiện chất lượng câu trả lời cuối cùng thêm hơn 12 điểm
• Giá trị của đội đỏ được xác nhận:Các nhà nghiên cứu bảo mật báo cáo rằng các dấu vết tư duy là "vô giá" để hiểu suy luận khai thác
• GRPO vượt trội hơn RLHF:Điểm mạch lạc tốt hơn 15% so với phương pháp truyền thống