ہم نے 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 سراغ فراہم کیے۔

ڈیپ سیک ڈیٹاسیٹ کی تشکیل

ٹریس جمع کرنے کا عمل

ہم نے متنوع ڈومینز سے سوچنے کے ٹریسز جمع کیے تاکہ جامع استدلال کی کوریج کو یقینی بنایا جا سکے:

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% بہتر ہم آہنگی کے سکور