Масштабирование вычислений LLM с использованием скрытых рассуждений: метод с рекуррентной глубиной

Экспериментальная архитектура LLM, которая способна масштабировать вычисления за счет скрытых рассуждений в латентном пространстве путем итеративного применения рекуррентного блока, что дает возможность развернуть вычисления на произвольную глубину.

Этот метод отличается от традиционных, которые увеличивают вычислительные ресурсы за счет генерации большего количества токенов. Например, в отличие от CoT, предложенный подход не требует специализированных датасетов, работает с небольшими окнами контекста и способен захватывать типы рассуждений, которые сложно выразить словами. В дополнение, модели этой архитектуры требуют меньше памяти для обучения и инференса.

Тестовая модель Huginn-3.5B (https://huggingface.co/tomg-group-umd/huginn-0125) получила 3.5 млрд параметров и была обучена на 800 млрд. токенов (веб-страницы, научные публикации и программный код) с использованием случайного числа итераций рекуррентного блока для каждой входной последовательности. Чтобы сократить потребление памяти использовалось усеченное обратное распространение, при котором градиенты вычисляются только для последних итераций.

Модель состоит из 3 основных блоков: прелюдии, рекуррентного блока и коды. Прелюдия преобразует входные данные в латентное пространство, рекуррентный блок выполняет итеративные вычисления, а кода преобразует латентное состояние обратно в вероятности токенов. Рекуррентный блок может быть повторен произвольное количество раз, позволяя модели выполнять произвольное количество вычислений перед генерацией токена.

Результаты проведенных тестов на стандартных задачах ARC, HellaSwag, MMLU свидетельствуют, что Huginn-3.5B превосходит традиционные модели на задачах, требующих сложных рассуждений (математические задачи и программирование). Например, на задачах GSM8k и MATH модель показала значительное улучшение производительности при увеличении числа рекуррентных итераций.

Модель не подвергалась файнтюну или посттренингу, но благодаря включению instruct-данных во время претрейна, она изначально понимает свой шаблон (https://huggingface.co/tomg-group-umd/huginn-0125#chat-templating) чата.

Чекпоинт на HF обучался всего на 47000 шагах и является академическим проектом.

Локальный инференс:

# Load the model  
 import torch 
 
 from transformers import AutoModelForCausalLM, AutoTokenizer 
 
 model = AutoModelForCausalLM.from_pretrained("tomg-group-umd/huginn-0125", torch_dtype=torch.bfloat16, trust_remote_code=True) 
 
 tokenizer = AutoTokenizer.from_pretrained("tomg-group-umd/huginn-0125") 
 
 # Modifying the Model's Depth at Test Time 
 
 input_ids = tokenizer.encode("The capital of Westphalia is", return_tensors="pt", add_special_tokens=True).to(device) 
 
 model.eval() 
 
 model.to(device) 
 
 model(input_ids, num_steps=32) 
 
 # Model can be used like a normal HF model 
 
 # You can provide `num_steps` directly to the `generate` call 
 
 model.eval() 
 
 config = GenerationConfig(max_length=256, stop_strings=["<|end_text|>", "<|end_turn|>"],  
 
                           use_cache=True, 
 
                           do_sample=False, temperature=None, top_k=None, top_p=None, min_p=None,  
 
                           return_dict_in_generate=True, 
 
                           eos_token_id=65505,bos_token_id=65504,pad_token_id=65509) 
 
 input_ids = tokenizer.encode("The capital of Westphalia is", return_tensors="pt", add_special_tokens=True).to(device) 
 
 outputs = model.generate(input_ids, config, tokenizer=tokenizer, num_steps=16)

Лицензирование: Apache 2.0 License.

Модель (https://huggingface.co/tomg-group-umd/huginn-0125)

Arxiv (https://arxiv.org/pdf/2502.05171)

GitHub (https://github.com/seal-rg/recurrent-pretraining)

Источник: github.com

Комментарии: