Clarification about Hyperparameters and CoLar stability

[mohammed-elkomy]

Hi @AlbertTan404

Thanks for the awesome work. I reviewed the code and the different training settings, and I trained a Llama 1b-instruct using CoT to replicate some CoLar results. I used GSM8K datasets; my CoT model has around 40% val acc.

However, I tried to use the math Colar SFT parameters for math you shared here.
https://huggingface.co/AlbertTan/CoLaR/blob/main/logs/colar/qsa-math/colar-math-sft-llama/hparams.yaml

Unfortunately, CoLar is unstable during training and stays at random guess performance even after 10 epochs.

Could you provide some guidelines on setting the parameters of the latent head and loss relative weights?

These are my current parameters.
latent_cot_config:
ce_weight: 1
embed_modeling_weight: 1
embed_modeling_loss: mse # {nll, mse}
entropy_weight: -1e-6
pred_embed_forward_weight: 1
max_compression_factor: 5
pred_compressed_cot: True
sqrt_mean: True

GSM8k
CoT checkpoint as initialization.

Thank you in advance.

[AlbertTan404]

Hi @mohammed-elkomy,

Could you please provide the loss/acc plot and the logged validation generations?
The config is good so I suppose there might be some bugs during training.

There are some of loss plots of mine: #5 (comment)

[AlbertTan404]

@mohammed-elkomy And also some samples from train/val data. Data formatting is crucial for training.

[mohammed-elkomy]

@AlbertTan404 Thanks a lot for the quick reply. It seems it's not a configuration issue. I tried a toy dataset of 32 samples and used it for train/test/val (we should expect a perfect performance). However, I got around 3% performance on Colar. Cot is fine-tuning correctly, but other latent models are not.

I will investigate this a bit more, but most probably it's a numerical stability issue or a bug. FYI, I installed a slightly higher version of PyTorch and Python. Is downgrading crucial, as you mentioned in Readme.md?

APPRECIATED!

[AlbertTan404]

Hi @mohammed-elkomy ,
I suggest that you could inspect the training from these perspectives (from the most important to less):

1. Checking whether the input text is in correct format (user prompt + latent generation)
2. Downgrading numpy < 2.0 since we have observed severe numerical instability with np>2.0 (not sure whether np is the root cause). For transformers/pytorch, we do not expect version incompatibility (If you find that it's a pytorch/transoformers version issue, please tell us)
3. Using pred_embed_forward_weight=0 (reduces training complexity)
4. Using max_compression_factor=2 or even 1 (helps convergence)
5. Using entropy_weight=0 (vulnerable to numerical instability)

The Latent Reasoning train&val process could be noisy in the initial stage, since there's a large domain gap between input embeddings and predicted latents. Thus, it requires a full-dataset training for at least 5 epochs (since we randomly sample c=1 to c=5 at each training step) to see the convergence.

If you have any updates or the issue remains, feel free to contact!

[mohammed-elkomy]

I tried downgrading numpy and am still experimenting with a toy dataset to test convergence. Generally, if the code is numerically stable, it should get a high score.

CoT => 99% for the toy gsm8k
Coconut => 40% for the toy gsm8k
CoLaR => (3%) still producing random tokens, although the loss got decreased; please check the attached screenshot. Probably the latent generation logic has some bug.

Some additional details: Here is a sample output of groundtruth and prediction by CoLaR after 90 epochs; it's producing text rather than te answer.

Is this a meaningful behavior?

Comparing answers: gt=5, pred=let h be the number of hours it takes joy to read 120 pages                                                                                                                          
Comparing answers: gt=16, pred=roque walks to and from work 3 times a week and each                                                                                                                                
Comparing answers: gt=16, pred=roque walks to work for 2 x 3 = <<2*                                                                                                                                                
Comparing answers: gt=7, pred=bobbie’s name is j-a-m-e-g-r-y so her                                                                                                                                                
Comparing answers: gt=7, pred=s then j for short for samantha then l for full name.                                

All predictions are generating max_new_tokens

Thanks so much for your time. This will be helpful for future collaborators interested in your awesome work.

[AlbertTan404]

Hi @mohammed-elkomy. Based on your training log, the root cause is underfitting. It would take more steps for converging CoLaR than CoT or Coconut, since we train an extra randomly initialized module (the latent head).

This repo has been used to reproduced CoLaR (see #5 (comment)). So I would suggest a full-dataset run on the gsm dataset till you see the real convergence of ce_loss (<1) and embed_modeling_loss (~0).

If performance issues persist after a full-dataset run, I'm open to an online meeting in Chinese or English :)