GRPO를 사용하여 Shannon V1.5가 생각하도록 훈련시킨 방법

1. 동기: 왜 사고 모델인가?

에 있어서,AI 레드팀연구에서, 이해하는 것은어떻게AI가 잠재적으로 유해한 결과에 도달하는지는 결과 자체를 연구하는 것만큼 중요합니다. Shannon V1.5의 사고 모델은 전체 사고 과정을 노출하여 연구자들이 다음을 수행할 수 있도록 합니다:

• 으로 이어지는 추론 경로를 추적하고검열되지 않은 AI의 결과적인행동
• 가드레일이 일반적으로 개입할 지점을 식별하며 (AI 가드레일의 중요성)
• 특정 추론 패턴을 목표로 하는 더 강력한 안전 메커니즘을 개발하고
• 관찰된 추론 체인을 기반으로 더 나은 적대적 테스트 사례를 생성합니다

2. GRPO 이해하기

그룹 상대 정책 최적화 (GRPO)은 전통적인 RLHF 방법론을 발전시킨 것으로, 추론 능력의 더 안정적이고 효율적인 훈련을 가능하게 합니다. DeepSeek AI가 개발했으며, 사고 과정 훈련에 특히 효과적인 것으로 입증되었습니다.

전통적인 RLHF보다 GRPO를 사용하는 이유?

GRPO 수학적 기초

GRPO는 절대적인 보상 모델과 비교하는 대신 그룹 내 응답을 비교하여 정책을 최적화합니다:

이러한 상대적 비교는 여러 가지 장점이 있습니다:

• 정규화:프롬프트별 다양한 난이도에 자동으로 조정
• 안정성:경사 추정치의 분산 감소
• 효율성:별도의 보상 모델 불필요

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. DeepSeek 증류

Shannon V1.5의 사고 능력을 부트스트랩하기 위해, DeepSeek의 추론 모델에서 사고 과정 패턴을 증류했습니다. 이는 우리의 사고 헤드를 훈련시키기 위한 고품질 CoT 추적을 제공했습니다.

DeepSeek 데이터셋 구성

추론 흔적 수집 과정

포괄적인 추론 범위를 보장하기 위해 다양한 도메인에서 사고 흔적을 수집했습니다:

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. 사고 헤드 아키텍처

Shannon V1.5 모델은 전용사고 헤드를 통합하여 최종 출력 전에 명시적인 추론 흔적을 생성합니다. 이 아키텍처 추가는 기본 Mixtral 아키텍처를 수정하지 않고도 투명한 CoT를 가능하게 합니다.

사고 헤드 구현

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. 훈련 파이프라인

1단계: 사고 헤드 사전 훈련

먼저, 표준 교차 엔트로피 손실을 사용하여 DeepSeek에서 증류된 CoT 추론 흔적에 대해 사고 헤드를 사전 훈련합니다:

# 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

2단계: GRPO 미세 조정

사전 훈련 후, 그룹 상대적 비교를 사용하여 사고 품질을 개선하기 위해 GRPO를 적용합니다:

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()
        }

3단계: 레드팀 전문화

마지막으로, 적대적 시나리오에 대해 추가로 튜닝하여 사고 흔적이 다음을 위한 추론을 적절하게 노출하도록 합니다:검열되지 않은 AI 결과분석:

6. 결과 및 분석

사고 품질 지표

주요 발견

• 투명성이 극적으로 향상됨:추론의 12%에서 97.8%가 이제 명시적으로 언어화됨
• 추론 정확도 증가:명시적 사고가 최종 답변 품질을 12점 이상 향상시킴
• 레드팀 가치 확인:보안 연구원들은 사고 흔적이 익스플로잇 추론을 이해하는 데 "매우 귀중하다"고 보고합니다.
• GRPO가 RLHF를 능가함:기존 접근 방식 대비 15% 더 나은 일관성 점수