LongRoPE
1.0.0
| etiquetas | conjuntos de datos | métrica | ||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
|
El artículo presenta LongRoPE, un método para ampliar la ventana de contexto de modelos de lenguaje grandes (LLM) más allá de los 2 millones de tokens.
Las ideas clave son:
Identifique y aproveche dos formas de no uniformidad en incrustaciones posicionales para minimizar la pérdida de información durante la interpolación. Esto permite una extensión de contexto 8x sin ajustes.
Utilice una estrategia de extensión progresiva eficiente con ajuste fino de 256k para alcanzar un contexto de 2048k, en lugar de ajustar directamente un contexto extremadamente grande.
Ajuste las incrustaciones para contextos más cortos para recuperar el rendimiento dentro del tamaño de ventana original.
El método se aplica a LLaMA2 y Mistral. Los experimentos en varias tareas demuestran la eficacia de LongRoPE para mantener el rendimiento en longitudes de contexto de 4k a 2048k.
La arquitectura Transformer lucha con la complejidad computacional cuadrática de la autoatención y su falta de generalización a posiciones simbólicas que no se ven en el momento del entrenamiento. Para escalar el cálculo de la autoatención a un contexto amplio, se han propuesto varios métodos, como RoPE, AliBi, sumideros de atención, etc. Sin embargo, ninguna de estas soluciones puede escalar efectivamente al contexto con millones de tokens y al mismo tiempo preservar la precisión del modelo. .
Este artículo presenta una nueva técnica, LongRoPE, que amplía la ventana de contexto de los LLM a más de 2 millones de tokens.
LongRoPE utiliza una estrategia de extensión progresiva para lograr una ventana de contexto de 2048k sin necesidad de realizar ajustes directos en textos excesivamente largos, que son raros y difíciles de conseguir. Esta estrategia comienza con una extensión de 256k en un LLM previamente entrenado, seguida de un ajuste fino a esta longitud.
Para abordar posibles caídas de rendimiento en la ventana de contexto original (más corta), LongRoPE ajusta aún más los factores de reescalado de RoPE en el LLM extendido, reduciendo a ventanas de contexto de 4k y 8k en el LLM ajustado de 256k utilizando su algoritmo de búsqueda para minimizar la interpolación posicional. Durante la inferencia de secuencias de menos de 8k de longitud, RoPE se actualiza con estos factores de reescala meticulosamente buscados.
Las pruebas en varios LLM y tareas que requieren contextos prolongados han validado la eficacia de LongRoPE. El método mantiene significativamente una baja perplejidad en longitudes de evaluación de 4k a 2048k tokens, logra una precisión superior al 90% en la recuperación de claves de acceso y ofrece una precisión comparable a los puntos de referencia estándar dentro de una ventana de contexto 4096.
Una mirada en profundidad a las modificaciones estructurales y sus implicaciones para el desempeño del modelo.
La arquitectura del modelo LongRoPE está diseñada para ampliar la ventana de contexto de los modelos de lenguaje grandes (LLM) a más de 2 millones de tokens, abordando las limitaciones de las arquitecturas Transformer tradicionales. La innovación clave reside en la estrategia de extensión progresiva y el ajuste de las incrustaciones posicionales.
Los componentes clave incluyen:
class RoPEPositionalEncoding ( nn . Module ):
def __init__ ( self , d_model , max_len = 1000000 , base = 10000 ):
super (). __init__ ()
self . d_model = d_model
self . max_len = max_len
self . base = base
self . theta = torch . tensor ([ base ** ( - 2 * ( i // 2 ) / d_model ) for i in range ( d_model )])
def forward ( self , positions ):
angles = positions . unsqueeze ( - 1 ) * self . theta
sin_cos = torch . stack ([ angles . cos (), angles . sin ()], dim = - 1 )
return sin_cos . view ( * sin_cos . shape [: - 2 ], - 1 ) def non_uniform_interpolation ( pos_embed , extension_ratio , lambda_factors , n_hat ):
d_model = pos_embed . shape [ - 1 ]
interpolated_pos = pos_embed . clone ()
for i in range ( d_model // 2 ):
mask = torch . arange ( pos_embed . shape [ - 2 ], device = pos_embed . device ) < n_hat
scale = torch . where ( mask , torch . ones_like ( pos_embed [..., 0 ], device = pos_embed . device ),
1 / ( lambda_factors [ i ] * extension_ratio ))
interpolated_pos [..., 2 * i ] *= scale
interpolated_pos [..., 2 * i + 1 ] *= scale
return interpolated_pos def progressive_extension ( model , data , base_length , target_length , population_size , num_mutations , num_crossovers , max_iterations ):
# Extend to 128k
lambda_factors_128k , n_hat_128k = search_lambda_factors ( model , data , 128000 / base_length , population_size , num_mutations , num_crossovers , max_iterations )
model = fine_tune ( model , data , 128000 , lambda_factors_128k , n_hat_128k , steps = 400 )
# Extend to 256k
lambda_factors_256k , n_hat_256k = search_lambda_factors ( model , data , 256000 / base_length , population_size , num_mutations , num_crossovers , max_iterations )
model = fine_tune ( model , data , 256000 , lambda_factors_256k , n_hat_256k , steps = 600 )
# Extend to target length
if target_length > 256000 :
final_lambda_factors , final_n_hat = search_lambda_factors ( model , data , target_length / base_length , population_size // 2 , num_mutations // 2 , num_crossovers // 2 , max_iterations // 2 )
model . lambda_factors [ "2048k" ] = final_lambda_factors
model . n_hat [ "2048k" ] = final_n_hat
return model , final_lambda_factors , final_n_hat , lambda_factors_256k , n_hat_256kLa arquitectura comienza con un LLM previamente entrenado y amplía su ventana de contexto de forma incremental. Inicialmente, el modelo está ajustado para manejar una longitud de contexto de 256k tokens. Este enfoque progresivo evita la necesidad de realizar ajustes directos en textos extremadamente largos, que son raros y computacionalmente costosos de procesar. Al aumentar gradualmente la longitud del contexto, el modelo puede adaptarse más eficazmente a secuencias más largas.
Para mantener el rendimiento en contextos de diferentes longitudes, LongRoPE ajusta las incrustaciones posicionales rotativas (RoPE). El modelo identifica y explota las no uniformidades en las incrustaciones posicionales para minimizar la pérdida de información durante la interpolación. Esto permite una extensión de contexto 8x sin necesidad de realizar ajustes. Además, el modelo emplea un algoritmo de búsqueda para encontrar factores de reescalado óptimos para contextos más cortos (por ejemplo, tokens de 4k y 8k) en el LLM ajustado de 256k. Estos ajustes garantizan que el modelo conserve un alto rendimiento incluso dentro del tamaño de la ventana de contexto original.
La arquitectura incorpora varias modificaciones estructurales para manejar eficientemente la mayor longitud del contexto:
Escala de capa : se realizan ajustes en la escala de las capas para garantizar la estabilidad y el rendimiento a medida que crece la ventana de contexto.
Gestión de la memoria : se emplean técnicas eficientes de gestión de la memoria para manejar contextos de gran tamaño sin abrumar los recursos del sistema.
Mecanismos de atención : se integran mecanismos de atención mejorados para garantizar que el modelo pueda centrarse en partes relevantes de la secuencia de entrada, incluso con el contexto extendido.
Atención por token : se introducen mecanismos de atención por token para capturar las relaciones contextuales entre los tokens, lo que permite que el modelo comprenda mejor el significado semántico de la entrada.
Los experimentos demuestran que LongRoPE mantiene una baja perplejidad en longitudes de evaluación de 4k a 2048k tokens y logra una alta precisión en tareas que requieren contextos largos. Esto lo hace adecuado para diversas aplicaciones, incluido el aprendizaje en contexto, el resumen de documentos extensos y el aprendizaje en pocas oportunidades.
Para obtener información más detallada, consulte el documento completo aquí.
Información sobre la codificación y los detalles operativos que permiten la funcionalidad de LongRoPE. Esto puede incluir fragmentos o pseudocódigo que ilustren componentes clave.
Para obtener información más detallada, consulte el documento.
Ejemplos completos que demuestran cómo aprovechar LongRoPE para diversas aplicaciones, desde análisis de texto hasta generación de documentos extensos.
# Example usage
data_path = "path/to/your/dataset"
d_model = 512
n_heads = 8
num_layers = 6
base_length = 4096
target_length = 2048 * 1024
data = load_data ( data_path )
model = LongRoPEModel ( d_model , n_heads , num_layers , base_length )
model = model . extend_context ( data , target_length )
input_ids = torch . randn ( 2 , target_length , d_model )
output = model ( input_ids )
print ( output . shape ) # Expected shape: (batch_size, target_length, d_model) Entrenamiento de conjuntos de datos personalizados
Para entrenar en un conjunto de datos personalizado:
Ajuste de hiperparámetros El rendimiento de LongRoPE puede ser sensible a los hiperparámetros. Los parámetros clave para ajustar incluyen:
population_size , num_mutations y num_crossovers en la búsqueda del factor lambda Tasa de aprendizaje y parámetros del programador para ajustar gradient_accumulation_steps para la estabilidad del entrenamiento
Mi implementación de LongRoPE logra los siguientes resultados:
Perplejidad:
Precisión de recuperación de claves de acceso:
Exactitud:
Comparación con modelos de referencia:
@article { ding2024longrope ,
title = { LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens } ,
author = { Ding, Yiran and Zhang, Li Lyna and Zhang, Chengruidong and Xu, Yuanyuan and Shang, Ning and Xu, Jiahang and Yang, Fan and Yang, Mao } ,
journal = { arXiv preprint arXiv:2402.13753 } ,
year = { 2024 }
}Nota: Este repositorio es un trabajo en progreso y aún no está listo para su uso en producción. Consulte el documento para obtener más detalles.
