FlexOlmo: Open Language Models for Flexible Data Use

# Merge two expert models via linear interpolation
import torch

def merge_models(base_state_dict, expert_state_dicts, weights):
    """
    base_state_dict: shared initialization weights
    expert_state_dicts: list of expert module state dicts
    weights: interpolation coefficients (sum to 1.0)
    """
    merged = {}
    for key in base_state_dict:
        # Compute task vectors (delta from base)
        task_vectors = [
            expert[key] - base_state_dict[key]
            for expert in expert_state_dicts
        ]
        # Weighted sum of task vectors
        combined = sum(w * tv for w, tv in zip(weights, task_vectors))
        merged[key] = base_state_dict[key] + combined
    return merged

import torch
import torch.nn as nn

class FlexRouter(nn.Module):
    """Learned router that selects which expert modules to activate."""
    def __init__(self, hidden_dim, num_experts, top_k=2):
        super().__init__()
        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
        self.top_k = top_k

    def forward(self, x):
        # x: [batch, seq_len, hidden_dim]
        logits = self.gate(x)  # [batch, seq_len, num_experts]
        weights, indices = torch.topk(logits, self.top_k, dim=-1)
        weights = torch.softmax(weights, dim=-1)
        return weights, indices