text capability recovery: weight merging

[sprnjt]

Problem:
Adding visual capabilities to Tiny Aya consistently degrades its multilingual text performance on benchmarks like m-ArenaHard and GlobalMGSM. With only 3.35B parameters, Tiny Aya has less capacity redundancy than the 8B/32B models where cross-modal merging has previously been validated.

Too early to merge in main (since it is Phase 1 rn)

Solution:
Implement a weight-merging script that linearly interpolates the fine-tuned VLM's LLM backbone with the original text-only Tiny Aya Base weights.

This lead to the idea of Pareto merging concept (which I did earlier in a research project, basically finding the best sweep or the sweet spot)

merged = (1 - α) × original_weight  +  α × finetuned_weight

The merge ratio α ∈ {0.3, 0.4, 0.5, 0.6, 0.7} is a CLI parameter, enabling a sweep to find the Pareto-optimal point between visual grounding and text task recovery.

What the script does?

Loads the original Tiny Aya base model, and fine tuned VLM checkpoint

Very important: Interpolates only language_model.* parameters — the multi_modal_projector and vision_encoder pass through untouched from the fine-tuned checkpoint

Validates key and shape compatibility before merging
Saves merged_state.pt and optionally a full HF model dir (--save-hf)

Usage:

# Single merge
python scripts/merge_weights.py \
  --original  CohereLabs/tiny-aya-base \
  --finetuned ./checkpoints/tiny-aya-vision-ft \
  --alpha     0.5 \
  --output    ./merged/alpha_0.5

# Sweep (bash)
for alpha in 0.3 0.4 0.5 0.6 0.7; do
  python scripts/merge_weights.py \
    --original CohereLabs/tiny-aya-base \
    --finetuned ./checkpoints/tiny-aya-vision-ft \
    --alpha $alpha \
    --output ./merged/alpha_$alpha
done

[sanggusti]

#11

Do you save some logs of the run? how does it look like?

[sprnjt]

The merge would happen after the fine tuning. I figured out the merging process and will need to fix the code as per the changes.